moonwhaler
a286e8b012
Extracts shared image helpers from Reference Concat into _image_ops.py and adds a new Reference Grid node with columns, fit_mode (cover/ contain/stretch), cell_size_source, fill_direction, and last_row_align controls. Reference Concat's legacy grid fallback is unchanged.
223 lines
9.5 KiB
Python
223 lines
9.5 KiB
Python
import math
|
|
|
|
import torch
|
|
|
|
from ._image_ops import FILL_PRESETS as _FILL_PRESETS
|
|
from ._image_ops import contain_fit as _contain_fit
|
|
from ._image_ops import cover_fit as _cover_fit
|
|
from ._image_ops import match_channels as _match_channels
|
|
from ._image_ops import normalize_refs as _normalize_refs
|
|
from ._image_ops import resize as _resize
|
|
|
|
CATEGORY = "MoonPack/image"
|
|
|
|
|
|
class ReferenceGrid:
|
|
DESCRIPTION = (
|
|
"Arranges a batch of reference images into a contact-sheet grid, with full "
|
|
"control over column count, per-cell fit mode, cell sizing, fill order and "
|
|
"alignment of an incomplete last row/column. Accepts either a MoonPack Image "
|
|
"List (keeps each reference's native resolution) or a plain stacked IMAGE batch."
|
|
)
|
|
SEARCH_ALIASES = ["grid", "contact sheet", "reference", "mosaic", "tile"]
|
|
INPUT_IS_LIST = True
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"reference_images": ("IMAGE", {
|
|
"tooltip": (
|
|
"Reference images to arrange into a grid. Connect a MoonPack Image List "
|
|
"to keep each image's native resolution (recommended for mixed-resolution "
|
|
"references), or a plain stacked IMAGE batch."
|
|
),
|
|
}),
|
|
"columns": ("INT", {
|
|
"default": 0, "min": 0, "max": 8192, "step": 1,
|
|
"tooltip": "Number of grid columns. 0 = auto (ceil(sqrt(n)), a roughly square grid).",
|
|
}),
|
|
"fit_mode": (["cover", "contain", "stretch"], {
|
|
"default": "cover",
|
|
"tooltip": (
|
|
"How each reference fills its cell. 'cover': scale to fill, center-crop "
|
|
"the overflow. 'contain': scale to fit inside, pad the leftover with fill "
|
|
"color. 'stretch': force an exact fit, distorting aspect ratio."
|
|
),
|
|
}),
|
|
"cell_size_source": (["largest", "smallest", "average"], {
|
|
"default": "largest",
|
|
"tooltip": (
|
|
"Which references set the cell size (per axis, independently): the "
|
|
"largest reference (upscales smaller ones), the smallest (downscales "
|
|
"larger ones), or the average dimensions."
|
|
),
|
|
}),
|
|
"cell_width": ("INT", {
|
|
"default": 0, "min": 0, "max": 8192, "step": 1,
|
|
"tooltip": "Fixed cell width override. 0 = use cell_size_source.",
|
|
}),
|
|
"cell_height": ("INT", {
|
|
"default": 0, "min": 0, "max": 8192, "step": 1,
|
|
"tooltip": "Fixed cell height override. 0 = use cell_size_source.",
|
|
}),
|
|
"fill_direction": (["row_major", "column_major"], {
|
|
"default": "row_major",
|
|
"tooltip": (
|
|
"Order references fill the grid. 'row_major': left to right, then top to "
|
|
"bottom. 'column_major': top to bottom, then left to right."
|
|
),
|
|
}),
|
|
"last_row_align": (["left", "center", "right"], {
|
|
"default": "left",
|
|
"tooltip": (
|
|
"Alignment of the incomplete trailing row (row_major) or column "
|
|
"(column_major, where left/right act as top/bottom) when the reference "
|
|
"count doesn't evenly divide the grid."
|
|
),
|
|
}),
|
|
"fill": (["black", "white", "gray", "edge_average"], {
|
|
"default": "black",
|
|
"tooltip": (
|
|
"Background color for empty cells and any 'contain' padding. "
|
|
"'edge_average' uses the mean color of all references."
|
|
),
|
|
}),
|
|
"interpolation": (["nearest", "bilinear", "bicubic", "lanczos"], {
|
|
"default": "bicubic",
|
|
"tooltip": "Resampling filter used whenever a reference image is resized.",
|
|
}),
|
|
"max_ref_side": ("INT", {
|
|
"default": 0, "min": 0, "max": 8192, "step": 1,
|
|
"tooltip": (
|
|
"Caps each reference's longest side to this many pixels before any other "
|
|
"processing (downscale only, never upscales). Keeps one oversized "
|
|
"reference from blowing up the grid size. 0 = no cap."
|
|
),
|
|
}),
|
|
"invert_mask": ("BOOLEAN", {
|
|
"default": False,
|
|
"tooltip": "Flip the output mask (empty/padding area becomes 1.0 instead of the reference area).",
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK")
|
|
RETURN_NAMES = ("image", "mask")
|
|
FUNCTION = "build"
|
|
CATEGORY = CATEGORY
|
|
|
|
def build(self, reference_images, columns, fit_mode, cell_size_source, cell_width, cell_height,
|
|
fill_direction, last_row_align, fill, interpolation, max_ref_side, invert_mask):
|
|
columns = columns[0]
|
|
fit_mode = fit_mode[0]
|
|
cell_size_source = cell_size_source[0]
|
|
cell_width = cell_width[0]
|
|
cell_height = cell_height[0]
|
|
fill_direction = fill_direction[0]
|
|
last_row_align = last_row_align[0]
|
|
fill = fill[0]
|
|
interpolation = interpolation[0]
|
|
max_ref_side = max_ref_side[0]
|
|
invert_mask = invert_mask[0]
|
|
|
|
refs = _normalize_refs(reference_images, max_ref_side, interpolation)
|
|
if not refs:
|
|
raise ValueError("ReferenceGrid: connect at least one reference image.")
|
|
|
|
n = len(refs)
|
|
channels = refs[0].shape[-1]
|
|
device, dtype = refs[0].device, refs[0].dtype
|
|
|
|
cols = columns if columns > 0 else max(1, math.ceil(math.sqrt(n)))
|
|
rows = math.ceil(n / cols)
|
|
|
|
heights = [ref.shape[1] for ref in refs]
|
|
widths = [ref.shape[2] for ref in refs]
|
|
if cell_size_source == "smallest":
|
|
src_h, src_w = min(heights), min(widths)
|
|
elif cell_size_source == "average":
|
|
src_h = round(sum(heights) / n)
|
|
src_w = round(sum(widths) / n)
|
|
else:
|
|
src_h, src_w = max(heights), max(widths)
|
|
cell_h = cell_height if cell_height > 0 else src_h
|
|
cell_w = cell_width if cell_width > 0 else src_w
|
|
|
|
if fill == "edge_average":
|
|
fill_rgb = torch.cat([ref.reshape(-1, ref.shape[-1]) for ref in refs], dim=0).mean(dim=0)
|
|
else:
|
|
fill_rgb = torch.tensor(_FILL_PRESETS[fill], device=device, dtype=dtype)
|
|
if channels != 3:
|
|
fill_rgb = fill_rgb.mean().expand(channels)
|
|
|
|
canvas_h, canvas_w = rows * cell_h, cols * cell_w
|
|
canvas = fill_rgb.view(1, 1, 1, channels).expand(1, canvas_h, canvas_w, channels).clone()
|
|
mask = torch.zeros((1, canvas_h, canvas_w), device=device, dtype=dtype)
|
|
|
|
positions = self._positions(n, cols, rows, fill_direction, last_row_align)
|
|
|
|
for ref, (r, c) in zip(refs, positions):
|
|
tile_ref = _match_channels(ref, channels)
|
|
y, x = r * cell_h, c * cell_w
|
|
if fit_mode == "cover":
|
|
tile = _cover_fit(tile_ref, cell_h, cell_w, interpolation)
|
|
canvas[:, y:y + cell_h, x:x + cell_w, :] = tile
|
|
mask[:, y:y + cell_h, x:x + cell_w] = 1.0
|
|
elif fit_mode == "contain":
|
|
tile, top, left, ch, cw = _contain_fit(tile_ref, cell_h, cell_w, interpolation, fill_rgb)
|
|
canvas[:, y:y + cell_h, x:x + cell_w, :] = tile
|
|
mask[:, y + top:y + top + ch, x + left:x + left + cw] = 1.0
|
|
else:
|
|
tile = _resize(tile_ref, cell_h, cell_w, interpolation)
|
|
canvas[:, y:y + cell_h, x:x + cell_w, :] = tile
|
|
mask[:, y:y + cell_h, x:x + cell_w] = 1.0
|
|
|
|
if invert_mask:
|
|
mask = 1.0 - mask
|
|
|
|
return (canvas, mask)
|
|
|
|
@staticmethod
|
|
def _positions(n, cols, rows, fill_direction, align):
|
|
"""Maps flat index i (0..n-1) to (row, col), honoring fill_direction and applying
|
|
last_row_align to whichever trailing group (last row, or last column for
|
|
column_major) is incomplete."""
|
|
if fill_direction == "row_major":
|
|
full_rows = n // cols
|
|
remainder = n - full_rows * cols
|
|
offset = 0
|
|
if remainder and align != "left":
|
|
gap = cols - remainder
|
|
offset = gap // 2 if align == "center" else gap
|
|
positions = []
|
|
for i in range(n):
|
|
r, c = divmod(i, cols)
|
|
if r == full_rows and remainder:
|
|
c += offset
|
|
positions.append((r, c))
|
|
return positions
|
|
else:
|
|
full_cols = n // rows
|
|
remainder = n - full_cols * rows
|
|
offset = 0
|
|
if remainder and align != "left":
|
|
gap = rows - remainder
|
|
offset = gap // 2 if align == "center" else gap
|
|
positions = []
|
|
for i in range(n):
|
|
c, r = divmod(i, rows)
|
|
if c == full_cols and remainder:
|
|
r += offset
|
|
positions.append((r, c))
|
|
return positions
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"MoonPack_ReferenceGrid": ReferenceGrid,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"MoonPack_ReferenceGrid": "Reference Grid",
|
|
}
|