Spaces:
Sleeping
Sleeping
| """ | |
| 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 | |