composer.py

"""
ReservedRegionFrameComposer — pure-Python port of the ComfyUI-BFSNodes node.

Adds a chroma-key side strip to every frame, placing the reference face inside
it so the LTX-2.3 model can use it as a persistent identity template throughout
generation.  After generation, call crop_reserved_region() to remove the strip.
"""

import math
from typing import Literal

import numpy as np
from PIL import Image


# ---------------------------------------------------------------------------
# Low-level helpers (ported from ComfyUI-BFSNodes/util.py)
# ---------------------------------------------------------------------------

def _fit_inside(src_w: int, src_h: int, max_w: int, max_h: int) -> tuple[int, int]:
    if src_w <= 0 or src_h <= 0:
        return 1, 1
    scale = min(max_w / src_w, max_h / src_h)
    return max(1, int(round(src_w * scale))), max(1, int(round(src_h * scale)))


def _aligned_offset(container: int, content: int, align: str) -> int:
    if align == "start":
        return 0
    if align == "end":
        return max(0, container - content)
    return max(0, (container - content) // 2)


def _paste_with_alpha(dst: Image.Image, src: Image.Image, xy: tuple[int, int]) -> None:
    if src.mode == "RGBA":
        dst.paste(src, xy, src.split()[-1])
    else:
        dst.paste(src, xy)


def _add_padding(img: Image.Image, pad: int = 16) -> Image.Image:
    canvas = Image.new("RGBA", (img.width + pad * 2, img.height + pad * 2), (255, 255, 255, 255))
    canvas.paste(img.convert("RGBA"), (pad, pad))
    return canvas


# ---------------------------------------------------------------------------
# Face layout helpers
# ---------------------------------------------------------------------------

def _layout_faces(
    faces: list[Image.Image],
    region_w: int,
    region_h: int,
    scale_pct: float,
    padding: int,
    gap: int,
    stack: str,
    align_main: str,
    align_cross: str,
) -> Image.Image:
    """Composite all faces into a single region_w x region_h RGBA tile."""
    canvas = Image.new("RGBA", (region_w, region_h), (0, 0, 0, 0))
    if not faces:
        return canvas

    n = len(faces)
    avail_w = region_w - 2 * padding
    avail_h = region_h - 2 * padding

    # Determine grid shape
    if stack == "horizontal" or (stack == "auto" and region_w >= region_h):
        cols, rows = n, 1
    elif stack == "vertical" or (stack == "auto" and region_h > region_w):
        cols, rows = 1, n
    else:  # grid
        cols = math.ceil(math.sqrt(n))
        rows = math.ceil(n / cols)

    cell_w = max(1, (avail_w - gap * (cols - 1)) // cols)
    cell_h = max(1, (avail_h - gap * (rows - 1)) // rows)

    for i, face in enumerate(faces):
        col, row = i % cols, i // cols
        fw, fh = _fit_inside(face.width, face.height, int(cell_w * scale_pct / 100), int(cell_h * scale_pct / 100))
        resized = face.resize((fw, fh), Image.LANCZOS)

        cx = padding + col * (cell_w + gap) + _aligned_offset(cell_w, fw, align_cross)
        cy = padding + row * (cell_h + gap) + _aligned_offset(cell_h, fh, align_main)
        _paste_with_alpha(canvas, resized, (cx, cy))

    return canvas


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def compose_frames(
    frames: np.ndarray,
    face_image: Image.Image,
    region_position: Literal["left", "right", "top", "bottom"] = "left",
    region_size_px: int = 256,
    face_scale_pct: float = 100.0,
    face_padding_px: int = 12,
    face_gap_px: int = 12,
    face_align_main: str = "center",
    face_align_cross: str = "center",
    chroma_rgb: tuple[int, int, int] = (0, 255, 0),
) -> np.ndarray:
    """
    Args:
        frames:         uint8 numpy array [N, H, W, 3]
        face_image:     PIL Image — the reference face
        region_*:       strip geometry and face placement
        chroma_rgb:     background colour of the reserved strip

    Returns:
        uint8 numpy array [N, H, W, 3] with the chroma strip composited in.
        The original content is shrunk to fill the remaining area; total
        resolution is unchanged (matches the input WxH).
    """
    N, H, W, C = frames.shape
    face_pil = _add_padding(face_image.convert("RGBA"), 16)

    vertical = region_position in ("top", "bottom")
    if vertical:
        content_h = H - region_size_px
        content_w = W
        region_w, region_h = W, region_size_px
    else:
        content_w = W - region_size_px
        content_h = H
        region_w, region_h = region_size_px, H

    # Pre-render the face tile (same for every frame)
    face_tile = _layout_faces(
        [face_pil],
        region_w, region_h,
        face_scale_pct, face_padding_px, face_gap_px,
        "auto", face_align_main, face_align_cross,
    )
    chroma_bg = Image.new("RGB", (region_w, region_h), chroma_rgb)
    chroma_bg.paste(face_tile, (0, 0), face_tile)  # alpha-composite face onto chroma

    out = np.empty_like(frames)

    for i in range(N):
        frame_pil = Image.fromarray(frames[i], "RGB")

        # Resize original content to fit the non-reserved area
        content_pil = frame_pil.resize((content_w, content_h), Image.LANCZOS)

        # Build full-size canvas
        canvas = Image.new("RGB", (W, H))
        if region_position == "left":
            canvas.paste(chroma_bg, (0, 0))
            canvas.paste(content_pil, (region_size_px, 0))
        elif region_position == "right":
            canvas.paste(content_pil, (0, 0))
            canvas.paste(chroma_bg, (content_w, 0))
        elif region_position == "top":
            canvas.paste(chroma_bg, (0, 0))
            canvas.paste(content_pil, (0, region_size_px))
        else:  # bottom
            canvas.paste(content_pil, (0, 0))
            canvas.paste(chroma_bg, (0, content_h))

        out[i] = np.array(canvas)

    return out


def crop_reserved_region(
    frames: np.ndarray,
    region_position: Literal["left", "right", "top", "bottom"] = "left",
    region_size_px: int = 256,
    output_size: tuple[int, int] | None = None,
) -> np.ndarray:
    """
    Remove the reserved strip from generated frames and resize back to
    output_size (W, H).  If output_size is None, resize to fill the full
    original frame dimensions.

    Args:
        frames:       uint8 [N, H, W, 3]
        output_size:  (W, H) to resize to after cropping, or None for original size

    Returns:
        uint8 [N, out_H, out_W, 3]
    """
    N, H, W, _ = frames.shape
    target_w = output_size[0] if output_size else W
    target_h = output_size[1] if output_size else H

    if region_position == "left":
        crop = frames[:, :, region_size_px:, :]
    elif region_position == "right":
        crop = frames[:, :, :W - region_size_px, :]
    elif region_position == "top":
        crop = frames[:, region_size_px:, :, :]
    else:  # bottom
        crop = frames[:, :H - region_size_px, :, :]

    if crop.shape[2] == target_w and crop.shape[1] == target_h:
        return crop

    out = np.empty((N, target_h, target_w, 3), dtype=np.uint8)
    for i in range(N):
        out[i] = np.array(Image.fromarray(crop[i]).resize((target_w, target_h), Image.LANCZOS))
    return out