moonwhaler-comfyui-moonpack/reference_concat_node.py

import math

import torch

from ._image_ops import FILL_PRESETS as _FILL_PRESETS
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 ReferenceConcat:
    DESCRIPTION = (
        "Concatenates a batch of reference images into a strip alongside a main image, "
        "for feeding identity references to video models without losing subject identity. "
        "The strip is placed on the main image's longest side, sized to exactly match that "
        "side's length (no resolution cap), and the main image is never shrunk. 'offset' "
        "slides the strip further into the main image (overlap) or further outside it (gap). "
        "If main_image is left unconnected, outputs a contact-sheet grid of just the "
        "reference images instead (side/ref_scale/offset are ignored in that mode). "
        "Accepts either a MoonPack Image List (keeps each reference's native resolution) or "
        "a plain stacked IMAGE batch."
    )
    SEARCH_ALIASES = ["concat", "reference", "identity", "first frame", "kjnodes", "video", "contact sheet"]
    INPUT_IS_LIST = True

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "reference_images": ("IMAGE", {
                    "tooltip": (
                        "Reference images to concatenate. Connect a MoonPack Image List to keep each "
                        "image's native resolution (recommended for mixed-resolution references), or a "
                        "plain stacked IMAGE batch."
                    ),
                }),
                "side": (["top", "bottom", "left", "right"], {
                    "default": "top",
                    "tooltip": (
                        "Which side of the main image the strip goes on. By default the strip's "
                        "row-vs-column layout is auto-picked from the main image's longest side, and "
                        "'top'/'left' place the strip before the main image along that axis while "
                        "'bottom'/'right' place it after. Enable force_side to make this side always "
                        "win instead. Ignored if main_image is unconnected."
                    ),
                }),
                "force_side": ("BOOLEAN", {
                    "default": False,
                    "tooltip": (
                        "Force the strip onto the exact side chosen above, overriding the automatic "
                        "row-vs-column pick: 'left'/'right' always stack the strip as a full-height "
                        "column, 'top'/'bottom' always stack it as a full-width row. The resulting strip "
                        "may end up thin if this fights the main image's orientation. Ignored if "
                        "main_image is unconnected."
                    ),
                }),
                "ref_scale": ("FLOAT", {
                    "default": 0.5, "min": 0.01, "max": 4.0, "step": 0.01,
                    "tooltip": (
                        "Strip thickness as a fraction of the main image's cross-side. The main image's "
                        "long side is a cap: the strip is only scaled down (never stretched) if it would "
                        "overflow. Ignored if main_image is unconnected."
                    ),
                }),
                "offset": ("INT", {
                    "default": 0, "min": -8192, "max": 8192, "step": 1,
                    "tooltip": (
                        "Positive: slide the strip into the main image, covering that many pixels of it. "
                        "Negative: pull the strip away from the main image, leaving a filled gap. "
                        "Clamped so the main image is never cropped or shrunk. Ignored if main_image is unconnected."
                    ),
                }),
                "fill": (["black", "white", "gray", "edge_average"], {
                    "default": "black",
                    "tooltip": (
                        "Background color for any gap/leftover space. 'edge_average' samples the main "
                        "image's edge nearest the strip, or the average of all references when main_image "
                        "is unconnected."
                    ),
                }),
                "interpolation": (["nearest", "bilinear", "bicubic", "lanczos"], {
                    "default": "bicubic",
                    "tooltip": (
                        "Resampling filter used whenever a reference image is resized, including "
                        "upscaling references smaller than the target size."
                    ),
                }),
                "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 contact-sheet or strip size. 0 = no cap."
                    ),
                }),
                "invert_mask": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "Flip the output mask (main-image/empty-cell area becomes 1.0 instead of the reference area).",
                }),
            },
            "optional": {
                "main_image": ("IMAGE", {
                    "tooltip": (
                        "Main / first-frame image. If a batch is passed, only the first frame is used. "
                        "Leave unconnected to instead output a contact-sheet grid of the reference images."
                    ),
                }),
            },
        }

    RETURN_TYPES = ("IMAGE", "MASK")
    RETURN_NAMES = ("image", "mask")
    FUNCTION = "concat"
    CATEGORY = CATEGORY

    def concat(self, reference_images, side, force_side, ref_scale, offset, fill, interpolation, max_ref_side,
               invert_mask, main_image=None):
        side = side[0]
        force_side = force_side[0]
        ref_scale = ref_scale[0]
        offset = offset[0]
        fill = fill[0]
        interpolation = interpolation[0]
        max_ref_side = max_ref_side[0]
        invert_mask = invert_mask[0]
        main = main_image[0][:1] if main_image is not None else None

        refs = _normalize_refs(reference_images, max_ref_side, interpolation)

        if main is None:
            if not refs:
                raise ValueError("ReferenceConcat: connect main_image, reference_images, or both.")
            return self._build_sheet(refs, fill, interpolation, invert_mask)

        _, main_h, main_w, channels = main.shape
        device, dtype = main.device, main.dtype

        if not refs:
            mask = torch.zeros((1, main_h, main_w), device=device, dtype=dtype)
            if invert_mask:
                mask = 1.0 - mask
            return (main, mask)

        if force_side:
            axis = "column" if side in ("left", "right") else "row"
        else:
            axis = "row" if main_w >= main_h else "column"
        before = side in ("top", "left")
        main_cross = main_h if axis == "row" else main_w
        target_length = main_w if axis == "row" else main_h

        nominal_thickness = max(1, round(ref_scale * main_cross))

        # length-per-unit-thickness for each ref: width/height for a row strip, height/width for a column strip.
        ratios = []
        for ref in refs:
            _, rh, rw, _ = ref.shape
            ratios.append((rw / rh) if axis == "row" else (rh / rw))

        if fill == "edge_average":
            edge_side = ("top" if axis == "row" else "left") if before else ("bottom" if axis == "row" else "right")
            if edge_side == "top":
                fill_rgb = main[0, 0, :, :].mean(dim=0)
            elif edge_side == "bottom":
                fill_rgb = main[0, -1, :, :].mean(dim=0)
            elif edge_side == "left":
                fill_rgb = main[0, :, 0, :].mean(dim=0)
            else:
                fill_rgb = main[0, :, -1, :].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)

        # ref_scale sets the nominal thickness directly. target_length is a cap, not a
        # forced fit: only scale down (never stretch) if the natural strip would overflow it.
        natural_total_length = sum(ratios) * nominal_thickness
        if natural_total_length > target_length:
            final_thickness = max(1, round(nominal_thickness * (target_length / natural_total_length)))
        else:
            final_thickness = nominal_thickness

        lengths = [max(1, round(final_thickness * ratio)) for ratio in ratios]
        strip_length = sum(lengths)
        if strip_length > target_length:
            # Rounding pushed a couple pixels past the cap; trim the last tile back to it.
            lengths[-1] = max(1, lengths[-1] - (strip_length - target_length))
            strip_length = sum(lengths)

        tiles = []
        for ref, length in zip(refs, lengths):
            ref = _match_channels(ref, channels)
            if axis == "row":
                tiles.append(_resize(ref, final_thickness, length, interpolation))
            else:
                tiles.append(_resize(ref, length, final_thickness, interpolation))

        strip = torch.cat(tiles, dim=2 if axis == "row" else 1)

        pad = target_length - strip_length
        if pad > 0:
            pad_shape = (1, final_thickness, pad, channels) if axis == "row" else (1, pad, final_thickness, channels)
            pad_tile = fill_rgb.view(1, 1, 1, channels).expand(pad_shape).clone()
            strip = torch.cat([strip, pad_tile], dim=2 if axis == "row" else 1)

        max_offset = min(final_thickness, main_cross)
        clamped_offset = min(offset, max_offset)

        if axis == "row":
            canvas_h = main_h + final_thickness - clamped_offset
            canvas_w = main_w
        else:
            canvas_h = main_h
            canvas_w = main_w + final_thickness - clamped_offset

        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)

        if axis == "row":
            if before:
                main_start, strip_start = final_thickness - clamped_offset, 0
            else:
                main_start, strip_start = 0, main_h - clamped_offset
            canvas[:, main_start:main_start + main_h, :, :] = main
            canvas[:, strip_start:strip_start + final_thickness, :, :] = strip
            mask[:, strip_start:strip_start + final_thickness, :] = 1.0
        else:
            if before:
                main_start, strip_start = final_thickness - clamped_offset, 0
            else:
                main_start, strip_start = 0, main_w - clamped_offset
            canvas[:, :, main_start:main_start + main_w, :] = main
            canvas[:, :, strip_start:strip_start + final_thickness, :] = strip
            mask[:, :, strip_start:strip_start + final_thickness] = 1.0

        if invert_mask:
            mask = 1.0 - mask

        return (canvas, mask)

    def _build_sheet(self, refs, fill, interpolation, invert_mask):
        # Justified-row layout (the algorithm behind Flickr/Google-Photos style contact
        # sheets): each row's images are scaled to a shared row height derived from their
        # own aspect ratios, so a wide reference gets a wide (or standalone) row and a tall
        # one gets a narrow row - never squeezed to match an unrelated reference's cell size.
        n = len(refs)
        channels = refs[0].shape[-1]
        device, dtype = refs[0].device, refs[0].dtype

        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)

        aspects = [ref.shape[2] / ref.shape[1] for ref in refs]
        avg_h = sum(ref.shape[1] for ref in refs) / n
        avg_w = sum(ref.shape[2] for ref in refs) / n
        cols = max(1, math.ceil(math.sqrt(n)))
        # Target width of a fully justified row, and the row height below which a row is
        # considered "full" - both derived from the references themselves, so the overall
        # canvas stays in the same ballpark as the old sqrt(n) grid without hardcoding a size.
        target_width = avg_w * cols
        target_row_height = avg_h

        row_groups = []
        current, aspect_sum = [], 0.0
        for i, a in enumerate(aspects):
            current.append(i)
            aspect_sum += a
            natural_height = target_width / aspect_sum
            if natural_height <= target_row_height:
                row_groups.append((current, natural_height))
                current, aspect_sum = [], 0.0
        if current:
            # Leftover row too sparse to fill target_width at target_row_height - use the
            # target height directly rather than stretching it taller to force-justify.
            row_groups.append((current, target_row_height))

        rows = []
        for group, height in row_groups:
            row_height = max(1, round(height))
            widths = [max(1, round(row_height * aspects[i])) for i in group]
            rows.append((group, row_height, widths))

        canvas_h = sum(row_height for _, row_height, _ in rows)
        canvas_w = max(sum(widths) for _, _, widths in rows)
        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)

        y = 0
        for group, row_height, widths in rows:
            x = 0
            for i, width in zip(group, widths):
                tile = _resize(_match_channels(refs[i], channels), row_height, width, interpolation)
                canvas[:, y:y + row_height, x:x + width, :] = tile
                mask[:, y:y + row_height, x:x + width] = 1.0
                x += width
            y += row_height

        if invert_mask:
            mask = 1.0 - mask

        return (canvas, mask)


NODE_CLASS_MAPPINGS = {
    "MoonPack_ReferenceConcat": ReferenceConcat,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "MoonPack_ReferenceConcat": "Reference Concat",
}