File size: 7,272 Bytes
1405c30 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 | """
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
|