custom_nodes/ComfyUI-BFSNodes/nodes.py

import math
import torch
from PIL import Image

from .util import (
    tensor_to_pil,
    pil_to_tensor,
    fit_inside,
    aligned_offset,
    paste_with_alpha,
    add_white_padding,
)


# ---------------------------------------------------------------------------
# FrameRangedFaceLoader
# ---------------------------------------------------------------------------

class FrameRangedFaceLoader:
    """
    Wraps a single face IMAGE with a frame range [frame_start, frame_end].
    Returns a FACE_SEQUENCE — a list of dicts used by ReservedRegionFrameComposer.

    frame_end = -1 means "until the last frame" (no upper limit).
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "frame_start": (
                    "INT",
                    {"default": 0, "min": 0, "max": 999999, "step": 1},
                ),
                "frame_end": (
                    "INT",
                    {"default": -1, "min": -1, "max": 999999, "step": 1},
                ),
            }
        }

    RETURN_TYPES = ("FACE_SEQUENCE",)
    RETURN_NAMES = ("face_sequence",)
    FUNCTION = "load"
    CATEGORY = "video/composition"

    def load(self, image, frame_start, frame_end):
        """
        image: IMAGE tensor [N, H, W, C] — only the first image is used.
        Returns a FACE_SEQUENCE list with a single entry.
        """
        face_pil = tensor_to_pil(image[0]).convert("RGBA")
        face_pil = add_white_padding(face_pil, 16)

        entry = {
            "image": face_pil,
            "frame_start": frame_start,
            "frame_end": frame_end,  # -1 == no upper limit
        }
        return ([entry],)


# ---------------------------------------------------------------------------
# FaceSequenceBatch
# ---------------------------------------------------------------------------

class FaceSequenceBatch:
    """
    Joins two FACE_SEQUENCE lists into one.
    Chain multiple nodes to build batches with many ranges.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "face_sequence_a": ("FACE_SEQUENCE",),
                "face_sequence_b": ("FACE_SEQUENCE",),
            }
        }

    RETURN_TYPES = ("FACE_SEQUENCE",)
    RETURN_NAMES = ("face_sequence",)
    FUNCTION = "batch"
    CATEGORY = "video/composition"

    def batch(self, face_sequence_a, face_sequence_b):
        return (face_sequence_a + face_sequence_b,)


# ---------------------------------------------------------------------------
# ReservedRegionFrameComposer
# ---------------------------------------------------------------------------

class ReservedRegionFrameComposer:
    """
    ComfyUI node that composes a reserved region (left/right/top/bottom)
    into every frame, while preserving the original output frame size.

    Main features:
    - Keeps final output resolution equal to original frame resolution
    - Fits video content into remaining area without crop
    - Fills reserved region with chroma color
    - Places one or many face images in that region
    - Supports temporal distribution modes for face batches (IMAGE input)
    - Supports per-frame ranges via FACE_SEQUENCE input
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "frames": ("IMAGE",),

                "region_position": (
                    ["left", "right", "top", "bottom"],
                    {"default": "left"}
                ),

                "region_size_px": (
                    "INT",
                    {"default": 320, "min": 8, "max": 8192, "step": 1}
                ),

                "face_distribution": (
                    [
                        "single_first",
                        "one_face_per_frame",
                        "one_face_per_interval",
                        "all_faces_every_frame"
                    ],
                    {"default": "one_face_per_interval"}
                ),

                "interval_frames": (
                    "INT",
                    {"default": 12, "min": 1, "max": 1000000, "step": 1}
                ),

                "overflow_mode": (
                    ["loop", "clamp", "error"],
                    {"default": "loop"}
                ),

                "stack_direction": (
                    ["auto", "vertical", "horizontal", "grid"],
                    {"default": "auto"}
                ),

                "face_scale_pct": (
                    "FLOAT",
                    {"default": 90.0, "min": 1.0, "max": 100.0, "step": 1.0}
                ),

                "face_padding_px": (
                    "INT",
                    {"default": 12, "min": 0, "max": 2048, "step": 1}
                ),

                "face_gap_px": (
                    "INT",
                    {"default": 12, "min": 0, "max": 2048, "step": 1}
                ),

                "face_align_main": (
                    ["start", "center", "end"],
                    {"default": "center"}
                ),

                "face_align_cross": (
                    ["start", "center", "end"],
                    {"default": "center"}
                ),

                "chroma_r": (
                    "INT",
                    {"default": 0, "min": 0, "max": 255, "step": 1}
                ),
                "chroma_g": (
                    "INT",
                    {"default": 255, "min": 0, "max": 255, "step": 1}
                ),
                "chroma_b": (
                    "INT",
                    {"default": 0, "min": 0, "max": 255, "step": 1}
                ),
            },
            "optional": {
                # When connected, overrides face_images completely.
                # face_distribution / interval_frames are ignored.
                "face_sequence": ("FACE_SEQUENCE",),

                # Legacy flat batch input (IMAGE). Used only when
                # face_sequence is NOT connected.
                "face_images": ("IMAGE",),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("frames_out",)
    FUNCTION = "process"
    CATEGORY = "video/composition"

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _prepare_face(self, face_tensor):
        """Converts input tensor to RGBA PIL so alpha transparency is preserved."""
        face = tensor_to_pil(face_tensor).convert("RGBA")
        face = add_white_padding(face, 16)
        return face

    # --- Resolution for plain IMAGE input ---

    def _resolve_faces_for_frame_image(
        self,
        faces_pil,
        frame_idx,
        face_distribution,
        interval_frames,
        overflow_mode,
    ):
        """
        Returns the list of face PIL images for the current frame
        when a plain IMAGE batch is used (legacy mode).
        """
        total = len(faces_pil)

        if total == 0:
            raise ValueError("No face images were provided.")

        if face_distribution == "single_first":
            return [faces_pil[0]]

        if face_distribution == "one_face_per_frame":
            idx = frame_idx
            if idx < total:
                return [faces_pil[idx]]
            if overflow_mode == "loop":
                return [faces_pil[idx % total]]
            elif overflow_mode == "clamp":
                return [faces_pil[-1]]
            else:
                raise ValueError(
                    f"Face batch is too small for one_face_per_frame. "
                    f"Frame index {frame_idx} requires face index {idx}, "
                    f"but only {total} face images were provided."
                )

        if face_distribution == "one_face_per_interval":
            idx = frame_idx // interval_frames
            if idx < total:
                return [faces_pil[idx]]
            if overflow_mode == "loop":
                return [faces_pil[idx % total]]
            elif overflow_mode == "clamp":
                return [faces_pil[-1]]
            else:
                raise ValueError(
                    f"Face batch is too small for one_face_per_interval. "
                    f"Frame index {frame_idx} requires interval face index {idx}, "
                    f"but only {total} face images were provided."
                )

        if face_distribution == "all_faces_every_frame":
            return faces_pil

        raise ValueError(f"Unsupported face_distribution mode: {face_distribution}")

    # --- Resolution for FACE_SEQUENCE input ---

    def _resolve_faces_for_frame_sequence(self, face_sequence, frame_idx, overflow_mode):
        """
        Returns the list of face PIL images for the current frame
        when a FACE_SEQUENCE is used (range mode).

        Collects all entries whose [frame_start, frame_end] covers frame_idx.
        frame_end == -1 means no upper limit.

        If no entry matches:
          - loop: wraps around using sorted order by frame_start
          - clamp: uses the last entry (highest frame_start)
          - error: raises ValueError
        """
        matched = [
            e for e in face_sequence
            if e["frame_start"] <= frame_idx and (
                e["frame_end"] == -1 or frame_idx <= e["frame_end"]
            )
        ]

        if matched:
            return [e["image"] for e in matched]

        # Fallback
        sorted_seq = sorted(face_sequence, key=lambda e: e["frame_start"])
        if not sorted_seq:
            raise ValueError("FACE_SEQUENCE is empty.")

        if overflow_mode == "loop":
            total = len(sorted_seq)
            return [sorted_seq[frame_idx % total]["image"]]
        elif overflow_mode == "clamp":
            return [sorted_seq[-1]["image"]]
        else:
            raise ValueError(
                f"No face entry covers frame {frame_idx} and overflow_mode is 'error'."
            )

    # --- Layout helpers ---

    def _resize_single_face(self, face, region_w, region_h, face_scale_pct, face_padding_px):
        """Resizes one face to fit inside the usable region area."""
        usable_w = max(1, region_w - 2 * face_padding_px)
        usable_h = max(1, region_h - 2 * face_padding_px)

        target_w = max(1, int(round(usable_w * (face_scale_pct / 100.0))))
        target_h = max(1, int(round(usable_h * (face_scale_pct / 100.0))))

        fw, fh = face.size
        tw, th = fit_inside(fw, fh, target_w, target_h)
        return face.resize((tw, th), Image.LANCZOS)

    def _layout_faces_stack(
        self,
        faces_pil,
        region_w,
        region_h,
        face_scale_pct,
        face_padding_px,
        face_gap_px,
        stack_direction,
    ):
        """
        Prepares resized face images for stack mode.
        Returns a tuple describing the chosen layout and resized images.
        """
        usable_w = max(1, region_w - 2 * face_padding_px)
        usable_h = max(1, region_h - 2 * face_padding_px)

        count = len(faces_pil)
        if count == 0:
            return None

        if stack_direction == "auto":
            stack_direction = "vertical" if region_h >= region_w else "horizontal"

        items = []

        if stack_direction == "vertical":
            slot_h = max(1, (usable_h - face_gap_px * (count - 1)) // count)
            slot_w = usable_w

            for face in faces_pil:
                fw, fh = face.size
                tw, th = fit_inside(
                    fw, fh,
                    max(1, int(round(slot_w * (face_scale_pct / 100.0)))),
                    max(1, int(round(slot_h * (face_scale_pct / 100.0))))
                )
                items.append(face.resize((tw, th), Image.LANCZOS))

            return ("vertical", items, usable_w, usable_h)

        if stack_direction == "horizontal":
            slot_w = max(1, (usable_w - face_gap_px * (count - 1)) // count)
            slot_h = usable_h

            for face in faces_pil:
                fw, fh = face.size
                tw, th = fit_inside(
                    fw, fh,
                    max(1, int(round(slot_w * (face_scale_pct / 100.0)))),
                    max(1, int(round(slot_h * (face_scale_pct / 100.0))))
                )
                items.append(face.resize((tw, th), Image.LANCZOS))

            return ("horizontal", items, usable_w, usable_h)

        if stack_direction == "grid":
            cols = max(1, math.ceil(math.sqrt(count)))
            rows = max(1, math.ceil(count / cols))

            cell_w = max(1, (usable_w - face_gap_px * (cols - 1)) // cols)
            cell_h = max(1, (usable_h - face_gap_px * (rows - 1)) // rows)

            for face in faces_pil:
                fw, fh = face.size
                tw, th = fit_inside(
                    fw, fh,
                    max(1, int(round(cell_w * (face_scale_pct / 100.0)))),
                    max(1, int(round(cell_h * (face_scale_pct / 100.0))))
                )
                items.append(face.resize((tw, th), Image.LANCZOS))

            return ("grid", items, usable_w, usable_h, cols, rows)

        raise ValueError(f"Unsupported stack_direction: {stack_direction}")

    def _paste_single_face(
        self,
        canvas,
        face,
        region_x,
        region_y,
        region_w,
        region_h,
        region_position,
        face_padding_px,
        face_align_main,
        face_align_cross,
    ):
        """Pastes one face inside the reserved region."""
        tw, th = face.size
        area_w = max(1, region_w - 2 * face_padding_px)
        area_h = max(1, region_h - 2 * face_padding_px)

        if region_position in ["left", "right"]:
            local_x = face_padding_px + aligned_offset(area_w, tw, face_align_cross)
            local_y = face_padding_px + aligned_offset(area_h, th, face_align_main)
        else:
            local_x = face_padding_px + aligned_offset(area_w, tw, face_align_main)
            local_y = face_padding_px + aligned_offset(area_h, th, face_align_cross)

        paste_with_alpha(canvas, face, (region_x + local_x, region_y + local_y))

    def _paste_stack_faces(
        self,
        canvas,
        faces_pil,
        region_x,
        region_y,
        region_w,
        region_h,
        face_scale_pct,
        face_padding_px,
        face_gap_px,
        face_align_main,
        face_align_cross,
        stack_direction,
    ):
        """Pastes multiple faces inside the reserved region."""
        layout = self._layout_faces_stack(
            faces_pil=faces_pil,
            region_w=region_w,
            region_h=region_h,
            face_scale_pct=face_scale_pct,
            face_padding_px=face_padding_px,
            face_gap_px=face_gap_px,
            stack_direction=stack_direction,
        )

        if layout is None:
            return

        usable_x = region_x + face_padding_px
        usable_y = region_y + face_padding_px

        kind = layout[0]

        if kind == "vertical":
            items   = layout[1]
            usable_w = layout[2]
            usable_h = layout[3]
            total_h = sum(img.size[1] for img in items) + face_gap_px * (len(items) - 1)
            start_y = usable_y + aligned_offset(usable_h, total_h, face_align_main)

            y = start_y
            for img in items:
                x = usable_x + aligned_offset(usable_w, img.size[0], face_align_cross)
                paste_with_alpha(canvas, img, (x, y))
                y += img.size[1] + face_gap_px
            return

        if kind == "horizontal":
            items    = layout[1]
            usable_w = layout[2]
            usable_h = layout[3]
            total_w = sum(img.size[0] for img in items) + face_gap_px * (len(items) - 1)
            start_x = usable_x + aligned_offset(usable_w, total_w, face_align_main)

            x = start_x
            for img in items:
                y = usable_y + aligned_offset(usable_h, img.size[1], face_align_cross)
                paste_with_alpha(canvas, img, (x, y))
                x += img.size[0] + face_gap_px
            return

        if kind == "grid":
            items    = layout[1]
            usable_w = layout[2]
            usable_h = layout[3]
            cols     = layout[4]  # type: ignore[index]
            rows     = layout[5]  # type: ignore[index]

            cell_w = max(1, (usable_w - face_gap_px * (cols - 1)) // cols)
            cell_h = max(1, (usable_h - face_gap_px * (rows - 1)) // rows)

            grid_w = cols * cell_w + (cols - 1) * face_gap_px
            grid_h = rows * cell_h + (rows - 1) * face_gap_px

            start_x = usable_x + aligned_offset(usable_w, grid_w, face_align_main)
            start_y = usable_y + aligned_offset(usable_h, grid_h, face_align_cross)

            for idx, img in enumerate(items):
                row = idx // cols
                col = idx % cols

                cell_x = start_x + col * (cell_w + face_gap_px)
                cell_y = start_y + row * (cell_h + face_gap_px)

                x = cell_x + (cell_w - img.size[0]) // 2
                y = cell_y + (cell_h - img.size[1]) // 2
                paste_with_alpha(canvas, img, (x, y))
            return

        raise ValueError(f"Unsupported stack layout kind: {kind}")

    # ------------------------------------------------------------------
    # Main execution
    # ------------------------------------------------------------------

    def process(
        self,
        frames,
        region_position,
        region_size_px,
        face_distribution,
        interval_frames,
        overflow_mode,
        stack_direction,
        face_scale_pct,
        face_padding_px,
        face_gap_px,
        face_align_main,
        face_align_cross,
        chroma_r,
        chroma_g,
        chroma_b,
        face_sequence=None,
        face_images=None,
    ):
        """Main node execution."""
        if len(frames.shape) != 4:
            raise ValueError(
                "Input 'frames' must be a batch IMAGE tensor with shape [N, H, W, C]."
            )

        # Determine input mode
        use_sequence = face_sequence is not None
        use_image = face_images is not None

        if not use_sequence and not use_image:
            raise ValueError(
                "You must connect either 'face_sequence' (FACE_SEQUENCE) "
                "or 'face_images' (IMAGE) to the node."
            )

        # Pre-process legacy IMAGE input
        if use_image and not use_sequence:
            if len(face_images.shape) != 4 or face_images.shape[0] < 1:
                raise ValueError(
                    "Input 'face_images' must be a valid IMAGE batch with at least one image."
                )
            face_count = face_images.shape[0]
            faces_pil_legacy = [self._prepare_face(face_images[i]) for i in range(face_count)]
        else:
            faces_pil_legacy = []

        n, orig_h, orig_w, _ = frames.shape

        if region_position in ["left", "right"]:
            max_region_size = max(1, orig_w - 1)
        else:
            max_region_size = max(1, orig_h - 1)

        region_size_px = max(1, min(region_size_px, max_region_size))

        if region_position in ["left", "right"]:
            region_w = region_size_px
            region_h = orig_h
            video_max_w = orig_w - region_size_px
            video_max_h = orig_h
        else:
            region_w = orig_w
            region_h = region_size_px
            video_max_w = orig_w
            video_max_h = orig_h - region_size_px

        if video_max_w < 1 or video_max_h < 1:
            raise ValueError(
                "The reserved region size is too large for the current frame resolution."
            )

        fitted_video_w, fitted_video_h = fit_inside(orig_w, orig_h, video_max_w, video_max_h)

        chroma_rgba = (chroma_r, chroma_g, chroma_b, 255)

        out_frames = []

        for i in range(n):
            frame_pil_src = tensor_to_pil(frames[i]).convert("RGB")
            frame_pil = frame_pil_src.resize((fitted_video_w, fitted_video_h), Image.LANCZOS)

            canvas = Image.new("RGBA", (orig_w, orig_h), color=(0, 0, 0, 255))

            if region_position == "left":
                region_x, region_y = 0, 0
                video_x, video_y = region_size_px, (orig_h - fitted_video_h) // 2

            elif region_position == "right":
                region_x, region_y = orig_w - region_size_px, 0
                video_x, video_y = 0, (orig_h - fitted_video_h) // 2

            elif region_position == "top":
                region_x, region_y = 0, 0
                video_x, video_y = (orig_w - fitted_video_w) // 2, region_size_px

            else:  # bottom
                region_x, region_y = 0, orig_h - region_size_px
                video_x, video_y = (orig_w - fitted_video_w) // 2, 0

            region_img = Image.new("RGBA", (region_w, region_h), color=chroma_rgba)
            canvas.paste(region_img, (region_x, region_y))
            canvas.paste(frame_pil.convert("RGBA"), (video_x, video_y))

            # Resolve faces for this frame
            if use_sequence:
                faces_for_frame = self._resolve_faces_for_frame_sequence(
                    face_sequence=face_sequence,
                    frame_idx=i,
                    overflow_mode=overflow_mode,
                )
            else:
                faces_for_frame = self._resolve_faces_for_frame_image(
                    faces_pil=faces_pil_legacy,
                    frame_idx=i,
                    face_distribution=face_distribution,
                    interval_frames=interval_frames,
                    overflow_mode=overflow_mode,
                )

            if len(faces_for_frame) == 1:
                single_face = self._resize_single_face(
                    face=faces_for_frame[0],
                    region_w=region_w,
                    region_h=region_h,
                    face_scale_pct=face_scale_pct,
                    face_padding_px=face_padding_px,
                )

                self._paste_single_face(
                    canvas=canvas,
                    face=single_face,
                    region_x=region_x,
                    region_y=region_y,
                    region_w=region_w,
                    region_h=region_h,
                    region_position=region_position,
                    face_padding_px=face_padding_px,
                    face_align_main=face_align_main,
                    face_align_cross=face_align_cross,
                )

            else:
                effective_stack_direction = stack_direction
                if stack_direction == "auto":
                    effective_stack_direction = (
                        "vertical" if region_position in ["left", "right"] else "horizontal"
                    )

                self._paste_stack_faces(
                    canvas=canvas,
                    faces_pil=faces_for_frame,
                    region_x=region_x,
                    region_y=region_y,
                    region_w=region_w,
                    region_h=region_h,
                    face_scale_pct=face_scale_pct,
                    face_padding_px=face_padding_px,
                    face_gap_px=face_gap_px,
                    face_align_main=face_align_main,
                    face_align_cross=face_align_cross,
                    stack_direction=effective_stack_direction,
                )

            out_frames.append(pil_to_tensor(canvas.convert("RGB")))

        out = torch.stack(out_frames, dim=0)
        return (out,)


# ---------------------------------------------------------------------------
# ComfyUI registration
# ---------------------------------------------------------------------------

NODE_CLASS_MAPPINGS = {
    "ReservedRegionFrameComposer": ReservedRegionFrameComposer,
    "FrameRangedFaceLoader": FrameRangedFaceLoader,
    "FaceSequenceBatch": FaceSequenceBatch,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ReservedRegionFrameComposer": "Reserved Region Frame Composer",
    "FrameRangedFaceLoader": "Frame Ranged Face Loader",
    "FaceSequenceBatch": "Face Sequence Batch",
}