diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..5dfa22e --- /dev/null +++ b/__init__.py @@ -0,0 +1,33 @@ +from .tiling import ImageTiler, ImageUntiler +from .transform import CropByCoords, PasteByCoords, ReorderBatch +from .coordinates import StringToIntegers, CombineCoords, SplitCoords + +NODE_CLASS_MAPPINGS = { + # From tiling.py + "ImageTiler": ImageTiler, + "ImageUntiler": ImageUntiler, + # From transform.py + "CropByCoords": CropByCoords, + "PasteByCoords": PasteByCoords, + "ReorderBatch": ReorderBatch, + # From coordinates.py + "StringToIntegers": StringToIntegers, + "CombineCoords": CombineCoords, + "SplitCoords": SplitCoords, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + # From tiling.py + "ImageTiler": "Tile Image", + "ImageUntiler": "Untile Image", + # From transform.py + "CropByCoords": "Crop Image by Coords", + "PasteByCoords": "Paste Image by Coords", + "ReorderBatch": "Reorder Batch", + # From coordinates.py + "StringToIntegers": "String to Integers", + "CombineCoords": "Combine Coords (to COORDS)", + "SplitCoords": "Split Coords (from COORDS)", +} + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/coordinates.py b/coordinates.py new file mode 100644 index 0000000..2438200 --- /dev/null +++ b/coordinates.py @@ -0,0 +1,66 @@ +from itertools import zip_longest + +def _parse_ints_from_string(s: str): + """Parses a comma-separated string into a list of integers, returning None on failure.""" + if not isinstance(s, str) or not s.strip(): + return None + try: + return [int(c.strip()) for c in s.split(',')] + except (ValueError, TypeError): + print(f"[ERROR] Could not parse coordinates string: '{s}'. Please use comma-separated integers.") + return None + +class StringToIntegers: + @classmethod + def INPUT_TYPES(cls): + return {"required": {"text": ("STRING", {"multiline": False, "default": ""})}} + + RETURN_TYPES = ("INTS",) + FUNCTION = "convert" + CATEGORY = "image/coordinates" + + def convert(self, text: str): + ints = _parse_ints_from_string(text) + # Provide a safe default for invalid or empty input to avoid breaking workflows. + return (ints if ints is not None else [0],) + +class CombineCoords: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "x_coords": ("INTS", {"default": [0]}), + "y_coords": ("INTS", {"default": [0]}), + } + } + + RETURN_TYPES = ("COORDS",) + FUNCTION = "combine" + CATEGORY = "image/coordinates" + + def combine(self, x_coords, y_coords): + # If one list is shorter, fill missing values with the last value of the respective list. + fill_x = x_coords[-1] if x_coords else 0 + fill_y = y_coords[-1] if y_coords else 0 + + combined = list(zip_longest(x_coords, y_coords, fillvalue=None)) + result = [[(x if x is not None else fill_x), (y if y is not None else fill_y)] for x, y in combined] + return (result,) + +class SplitCoords: + @classmethod + def INPUT_TYPES(cls): + return {"required": {"coords": ("COORDS",)}} + + RETURN_TYPES = ("INTS", "INTS") + RETURN_NAMES = ("x_coords", "y_coords") + FUNCTION = "split" + CATEGORY = "image/coordinates" + + def split(self, coords): + if not coords or not isinstance(coords, list): + return ([0], [0]) + + x_coords = [coord[0] for coord in coords] + y_coords = [coord[1] for coord in coords] + return (x_coords, y_coords) \ No newline at end of file diff --git a/tiling.py b/tiling.py new file mode 100644 index 0000000..76abe59 --- /dev/null +++ b/tiling.py @@ -0,0 +1,165 @@ +import torch +import torch.nn.functional as F + +class ImageTiler: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "rows": ("INT", {"default": 2, "min": 1, "max": 256}), + "cols": ("INT", {"default": 2, "min": 1, "max": 256}), + "overlap": ("INT", {"default": 64, "min": 0, "max": 1024, "step": 2}), + } + } + + RETURN_TYPES = ("IMAGE", "INT", "INT") + RETURN_NAMES = ("tiles", "original_width", "original_height") + FUNCTION = "execute" + CATEGORY = "image/tiling" + + def execute(self, image, rows, cols, overlap): + b, h, w, c = image.shape + original_h, original_w = h, w + + # Ensure overlap is an even number for symmetrical padding. + overlap = (overlap // 2) * 2 + padding = overlap // 2 + + # 1. Pad canvas to be perfectly divisible by rows/cols. + h_padded = (h + rows - 1) // rows * rows if rows > 0 else h + w_padded = (w + cols - 1) // cols * cols if cols > 0 else w + if h_padded != h or w_padded != w: + pad_h, pad_w = h_padded - h, w_padded - w + image = F.pad(image.permute(0, 3, 1, 2), (0, pad_w, 0, pad_h), mode='constant', value=0).permute(0, 2, 3, 1) + + # 2. Calculate tile dimensions including overlap. + base_tile_h = h_padded // rows + base_tile_w = w_padded // cols + + # Prevent overlap from being larger than the tile itself. + final_padding_y = min(padding, base_tile_h // 2) if rows > 1 else 0 + final_padding_x = min(padding, base_tile_w // 2) if cols > 1 else 0 + + output_tile_h = base_tile_h + (final_padding_y * 2) + output_tile_w = base_tile_w + (final_padding_x * 2) + + # 3. Slice image into tiles. + tiles = [] + with torch.no_grad(): + for img in image: + for i in range(rows): + for j in range(cols): + y1_src = max(0, i * base_tile_h - final_padding_y) + y2_src = min(h_padded, (i + 1) * base_tile_h + final_padding_y) + x1_src = max(0, j * base_tile_w - final_padding_x) + x2_src = min(w_padded, (j + 1) * base_tile_w + final_padding_x) + + source_slice = img[y1_src:y2_src, x1_src:x2_src, :] + + # Pad edge tiles to match the full tile size, ensuring a consistent batch. + # 'replicate' mode is used to avoid black borders on edge tiles. + pad_top = 0 if i > 0 else final_padding_y - (i * base_tile_h - y1_src) + pad_bottom = 0 if i < rows - 1 else final_padding_y - (y2_src - (i + 1) * base_tile_h) + pad_left = 0 if j > 0 else final_padding_x - (j * base_tile_w - x1_src) + pad_right = 0 if j < cols - 1 else final_padding_x - (x2_src - (j + 1) * base_tile_w) + + # Correction for single row/col cases + if rows == 1: pad_top = pad_bottom = 0 + if cols == 1: pad_left = pad_right = 0 + + padded_slice = F.pad(source_slice.permute(2, 0, 1).unsqueeze(0), + (pad_left, pad_right, pad_top, pad_bottom), mode='replicate') + + tiles.append(padded_slice.squeeze(0).permute(1, 2, 0).unsqueeze(0)) + + return (torch.cat(tiles, dim=0), original_w, original_h) + + +class ImageUntiler: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "tiles": ("IMAGE",), + "rows": ("INT", {"default": 2, "min": 1, "max": 256}), + "cols": ("INT", {"default": 2, "min": 1, "max": 256}), + "overlap": ("INT", {"default": 64, "min": 0, "max": 1024, "step": 2}), + }, + "optional": { + "original_width": ("INT", {"default": 0, "min": 0, "step": 1, "display": "number"}), + "original_height": ("INT", {"default": 0, "min": 0, "step": 1, "display": "number"}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "execute" + CATEGORY = "image/tiling" + + def execute(self, tiles, rows, cols, overlap, original_width=0, original_height=0): + group_size = rows * cols + if tiles.shape[0] % group_size != 0: + raise ValueError(f"Number of tiles ({tiles.shape[0]}) is not divisible by rows*cols ({group_size}).") + + overlap = (overlap // 2) * 2 + padding = overlap // 2 + + num_groups = tiles.shape[0] // group_size + output_images = [] + + with torch.no_grad(): + for g in range(num_groups): + group_tiles = tiles[g * group_size:(g + 1) * group_size] + _, tile_h_full, tile_w_full, c = group_tiles.shape + + base_tile_h = tile_h_full - overlap + base_tile_w = tile_w_full - overlap + + # 1. Prepare canvases for image data and blending weights. + final_h, final_w = rows * base_tile_h, cols * base_tile_w + canvas_h, canvas_w = final_h + overlap, final_w + overlap + image_canvas = torch.zeros((1, canvas_h, canvas_w, c), device=tiles.device, dtype=tiles.dtype) + weights_canvas = torch.zeros_like(image_canvas) + + # 2. Generate blending masks and assemble tiles. + for i in range(rows): + for j in range(cols): + # Create a linear fade mask (window) for blending. + fade = torch.linspace(0, 1, padding, device=tiles.device, dtype=tiles.dtype) + window_y = torch.ones(tile_h_full, device=tiles.device, dtype=tiles.dtype) + if padding > 0: + window_y[:padding], window_y[-padding:] = fade, fade.flip(0) + + window_x = torch.ones(tile_w_full, device=tiles.device, dtype=tiles.dtype) + if padding > 0: + window_x[:padding], window_x[-padding:] = fade, fade.flip(0) + + # At the final image edges, the mask should be solid, not faded. + if i == 0 and rows > 1: window_y[:padding] = 1.0 + if i == rows - 1 and rows > 1: window_y[-padding:] = 1.0 + if j == 0 and cols > 1: window_x[:padding] = 1.0 + if j == cols - 1 and cols > 1: window_x[-padding:] = 1.0 + + mask = (window_y.unsqueeze(1) * window_x.unsqueeze(0)).unsqueeze(0).unsqueeze(-1) + + # Add weighted tile to canvas. + idx = i * cols + j + tile = group_tiles[idx:idx+1] + y1, x1 = i * base_tile_h, j * base_tile_w + image_canvas[:, y1:y1+tile_h_full, x1:x1+tile_w_full, :] += tile * mask + weights_canvas[:, y1:y1+tile_h_full, x1:x1+tile_w_full, :] += mask + + # 3. Normalize canvas and crop to final dimensions. + # Use clamp to prevent division by zero in areas with no tile data. + final_canvas = image_canvas / weights_canvas.clamp(min=1e-6) + assembled_image = final_canvas[:, padding : padding + final_h, padding : padding + final_w, :] + + # Crop back to original dimensions if provided. + if original_width > 0 and original_height > 0: + final_image = assembled_image[:, :original_height, :original_width, :] + else: + final_image = assembled_image + + output_images.append(final_image) + + return (torch.cat(output_images, dim=0),) \ No newline at end of file diff --git a/transform.py b/transform.py new file mode 100644 index 0000000..67bf6c3 --- /dev/null +++ b/transform.py @@ -0,0 +1,197 @@ +import torch +import torch.nn.functional as F +from .coordinates import _parse_ints_from_string, CombineCoords + +class CropByCoords: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "width": ("INT", {"default": 64, "min": 1}), + "height": ("INT", {"default": 64, "min": 1}), + }, + "optional": { + "coords": ("COORDS",), + "x_coords_str": ("STRING", {"multiline": False}), + "y_coords_str": ("STRING", {"multiline": False}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "crop" + CATEGORY = "image/transform" + + def crop(self, images, width, height, coords=None, x_coords_str="", y_coords_str=""): + if coords is None and not x_coords_str and not y_coords_str: + return (images,) # No operation if no coordinates are provided + + if coords is None: + xs = _parse_ints_from_string(x_coords_str) or [0] + ys = _parse_ints_from_string(y_coords_str) or [0] + coords, _ = CombineCoords().combine(xs, ys) + + b, H, W, C = images.shape + + # Fast path: If all crops are identical, perform a single batch operation. + if len(coords) == 1: + x, y = coords[0] + x1_src, y1_src = max(0, x), max(0, y) + x2_src, y2_src = min(W, x + width), min(H, y + height) + + crop_slice = images[:, y1_src:y2_src, x1_src:x2_src, :] + slice_h, slice_w = crop_slice.shape[1:3] + + # Create a padded canvas to handle out-of-bounds crops. + padded_crop = torch.zeros((b, height, width, C), device=images.device, dtype=images.dtype) + paste_y_start, paste_x_start = max(0, -y), max(0, -x) + padded_crop[:, paste_y_start:paste_y_start+slice_h, paste_x_start:paste_x_start+slice_w, :] = crop_slice + return (padded_crop,) + + # Iterative path: Process each image with its corresponding coordinate. + crops = [] + for i in range(b): + # Use the last coordinate if the list is shorter than the batch size. + coord = coords[i if i < len(coords) else -1] + x, y = coord[0], coord[1] + + x1_src, y1_src = max(0, x), max(0, y) + x2_src, y2_src = min(W, x + width), min(H, y + height) + crop_slice = images[i:i+1, y1_src:y2_src, x1_src:x2_src, :] + slice_h, slice_w = crop_slice.shape[1:3] + + padded_crop = torch.zeros((1, height, width, C), device=images.device, dtype=images.dtype) + paste_y_start, paste_x_start = max(0, -y), max(0, -x) + padded_crop[:, paste_y_start:paste_y_start+slice_h, paste_x_start:paste_x_start+slice_w, :] = crop_slice + crops.append(padded_crop) + + return (torch.cat(crops, dim=0),) + + +class PasteByCoords: + @classmethod + def INPUT_TYPES(cls): + return { + "required": {"base": ("IMAGE",), "inserts": ("IMAGE",)}, + "optional": { + "coords": ("COORDS",), + "mask": ("MASK",), + "x_coords_str": ("STRING", {"multiline": False}), + "y_coords_str": ("STRING", {"multiline": False}), + "feathering": ("INT", {"default": 8, "min": 0, "max": 1024}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "paste" + CATEGORY = "image/transform" + + def paste(self, base, inserts, coords=None, mask=None, x_coords_str="", y_coords_str="", feathering=0): + if coords is None and not x_coords_str and not y_coords_str: + coords = [[0,0]] # Default to top-left corner if no coords provided. + + if coords is None: + xs = _parse_ints_from_string(x_coords_str) or [0] + ys = _parse_ints_from_string(y_coords_str) or [0] + coords, _ = CombineCoords().combine(xs, ys) + + b, H, W, C = base.shape + result = base.clone() + + num_inserts, num_coords = inserts.shape[0], len(coords) + + for i in range(b): + # Cycle through coords and inserts if their lists are shorter than the base batch. + coord = coords[i % num_coords] + x, y = coord[0], coord[1] + + ins = inserts[i % num_inserts].unsqueeze(0) + _, h, w, _ = ins.shape + + # --- Blending Mask Generation --- + if mask is not None: + m = mask[i % mask.shape[0]] + # Resize mask to match insert dimensions if necessary. + if m.shape[0] != h or m.shape[1] != w: + m = F.interpolate(m.unsqueeze(0).unsqueeze(0), size=(h, w), mode="bilinear", align_corners=False).squeeze(0).squeeze(0) + blending_mask = m.unsqueeze(0).unsqueeze(-1) + else: + blending_mask = torch.ones((1, h, w, 1), device=ins.device, dtype=ins.dtype) + # Apply feathering if specified, creating a soft edge. + if feathering > 0: + f = min(feathering, h // 2, w // 2) + if f > 0: + fade = torch.linspace(0, 1, f, device=ins.device, dtype=ins.dtype) + fade_x = fade.view(1, 1, -1, 1) + fade_y = fade.view(1, -1, 1, 1) + blending_mask[:, :f, :, :] *= fade_y + blending_mask[:, -f:, :, :] *= fade_y.flip(1) + blending_mask[:, :, :f, :] *= fade_x + blending_mask[:, :, -f:, :] *= fade_x.flip(2) + + blending_mask = blending_mask.repeat(1, 1, 1, C) + + # --- Paste Operation --- + y1_dst, x1_dst = max(0, y), max(0, x) + y2_dst, x2_dst = min(H, y + h), min(W, x + w) + y1_src, x1_src = max(0, -y), max(0, -x) + y2_src, x2_src = y1_src + (y2_dst - y1_dst), x1_src + (x2_dst - x1_dst) + + if y1_dst >= y2_dst or x1_dst >= x2_dst: continue + + # Alpha composite the insert over the base image. + base_crop = result[:, y1_dst:y2_dst, x1_dst:x2_dst, :] + mask_crop = blending_mask[:, y1_src:y2_src, x1_src:x2_src, :] + ins_crop = (ins * blending_mask)[:, y1_src:y2_src, x1_src:x2_src, :] + + result[:, y1_dst:y2_dst, x1_dst:x2_dst, :] = base_crop * (1 - mask_crop) + ins_crop + + return (result,) + +class ReorderBatch: + MODES = ["row_to_column", "reverse", "ping-pong"] + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE",), + "mode": (cls.MODES, ), + }, + "optional": { + "columns": ("INT", {"default": 2, "min": 1, "max": 128}), + } + } + + RETURN_TYPES = ("IMAGE", "INT") + RETURN_NAMES = ("images", "new_batch_size") + FUNCTION = "reorder" + CATEGORY = "image/transform" + + def reorder(self, images: torch.Tensor, mode: str, columns: int = 2): + batch_size = images.shape[0] + if batch_size == 0: + return (images, 0) + + output_images = images + + if mode == "reverse": + output_images = torch.flip(images, dims=[0]) + + elif mode == "ping-pong": + # Creates a seamless loop by appending the reversed sequence (excluding start/end frames). + if batch_size > 2: + return_trip = images[1:-1].flip(dims=[0]) + output_images = torch.cat((images, return_trip), dim=0) + elif batch_size == 2: # For 2 frames, just add the first one at the end. + output_images = torch.cat((images, images[0:1]), dim=0) + + elif mode == "row_to_column": + if columns <= 0: return (images, batch_size) + + # Transposes the batch order. E.g., for 2 cols: [0,1,2,3,4,5] -> [0,2,4,1,3,5] + rows = (batch_size + columns - 1) // columns + indices = [row * columns + col for col in range(columns) for row in range(rows) if row * columns + col < batch_size] + output_images = images[torch.tensor(indices)] + + return (output_images, output_images.shape[0]) \ No newline at end of file