import math import comfy.model_management import torch import torch.nn.functional as F from comfy_api.latest import io from kornia.geometry.transform.pyramid import ( PyrUp, build_laplacian_pyramid, build_pyramid, find_next_powerof_two, is_powerof_two, ) from torch import Tensor from .nodes_registry import comfy_node _CHUNK_SIZE = 8 _MASK_LOW_RES_LONG_SIDE = 64 def _pad_for_laplacian(image: torch.Tensor) -> tuple[torch.Tensor, tuple[int, int]]: """Pad an image tensor so H and W are powers of two (kornia requirement).""" h, w = image.shape[2], image.shape[3] pad_right = 0 pad_down = 0 if not (is_powerof_two(h) and is_powerof_two(w)): pad_right = find_next_powerof_two(w) - w pad_down = find_next_powerof_two(h) - h image = F.pad(image, (0, pad_right, 0, pad_down), mode="reflect") return image, (pad_right, pad_down) def _gaussian_pyramid( images: torch.Tensor, max_level: int, border_type: str = "reflect", align_corners: bool = False, ) -> list[Tensor]: h, w = images.shape[2], images.shape[3] if not (is_powerof_two(w) and is_powerof_two(h)): padding = (0, find_next_powerof_two(w) - w, 0, find_next_powerof_two(h) - h) images = F.pad(images, padding, mode=border_type) return build_pyramid(images, max_level, border_type, align_corners) def _resize_preserving_aspect_ratio( images: torch.Tensor, long_side: int, mode: str ) -> torch.Tensor: h, w = images.shape[-2:] current_long_side = max(h, w) if current_long_side == long_side: return images scale = long_side / current_long_side resized_h = max(1, int(round(h * scale))) resized_w = max(1, int(round(w * scale))) if mode == "nearest": return F.interpolate(images, size=(resized_h, resized_w), mode=mode) return F.interpolate( images, size=(resized_h, resized_w), mode=mode, align_corners=False, ) def _apply_low_res_mask_dilation( mask: torch.Tensor, spatial_radius: int, long_side: int = _MASK_LOW_RES_LONG_SIDE, ) -> torch.Tensor: if spatial_radius <= 0: return mask original_size = mask.shape[-2:] mask_low_res = _resize_preserving_aspect_ratio( mask.float(), long_side, mode="bilinear" ) mask_low_res = F.max_pool2d( mask_low_res, kernel_size=spatial_radius * 2 + 1, stride=1, padding=spatial_radius, ) return F.interpolate( mask_low_res, size=original_size, mode="bilinear", align_corners=False, ) def _pyramid_blend_chunk( image1: torch.Tensor, image2: torch.Tensor, mask: torch.Tensor, max_level: int = 7, ) -> torch.Tensor: """Blend a single chunk (already padded, already on device).""" pyramid1 = build_laplacian_pyramid(image1, max_level=max_level) pyramid2 = build_laplacian_pyramid(image2, max_level=max_level) pyramid_mask = _gaussian_pyramid(mask, max_level=max_level) pyr_up = PyrUp() output = pyramid1[-1] * pyramid_mask[-1] + pyramid2[-1] * (1 - pyramid_mask[-1]) for i in range(len(pyramid1) - 2, -1, -1): residual = pyramid1[i] * pyramid_mask[i] + pyramid2[i] * (1 - pyramid_mask[i]) output = pyr_up(output) + residual return output def _pyramid_blend( image1: torch.Tensor, image2: torch.Tensor, mask: torch.Tensor, max_level: int = 7, device: torch.device | None = None, output_device: torch.device | None = None, ) -> torch.Tensor: if image1.shape != image2.shape: raise ValueError( f"input images must have the same size, {image1.shape} != {image2.shape}" ) if image1.shape[0] != mask.shape[0]: raise ValueError( "image_a, image_b, and mask must have the same frame count for blending" ) if image1.shape[-2:] != mask.shape[-2:]: raise ValueError( "image_a, image_b, and mask must have the same spatial resolution for blending" ) _, padding = _pad_for_laplacian(image1[:1]) orig_h, orig_w = image1.shape[-2], image1.shape[-1] padded_min = min(orig_h + padding[1], orig_w + padding[0]) max_level = min(max_level, int(math.log2(padded_min))) if any(padding): mask = torch.nn.functional.pad( mask, (0, padding[0], 0, padding[1]), mode="reflect" ) B = image1.shape[0] results = [] for start in range(0, B, _CHUNK_SIZE): end = min(start + _CHUNK_SIZE, B) img1_chunk, _ = _pad_for_laplacian(image1[start:end]) img2_chunk, _ = _pad_for_laplacian(image2[start:end]) mask_chunk = mask[start:end] if device is not None: img1_chunk = img1_chunk.to(device) img2_chunk = img2_chunk.to(device) mask_chunk = mask_chunk.to(device) out_chunk = _pyramid_blend_chunk( img1_chunk, img2_chunk, mask_chunk, max_level=max_level ) cropped = out_chunk[..., :orig_h, :orig_w].clamp(0, 1) results.append( cropped.to(output_device) if output_device is not None else cropped ) return torch.cat(results, dim=0) @comfy_node( name="LTXVLaplacianPyramidBlend", description="LTX Laplacian Pyramid Blend", ) class LTXVLaplacianPyramidBlend(io.ComfyNode): """Blend two images seamlessly using Laplacian pyramid blending with a mask.""" @classmethod def define_schema(cls): return io.Schema( node_id="LTXVLaplacianPyramidBlend", category="Lightricks/utility", description="Blend two images seamlessly using Laplacian pyramid blending.", inputs=[ io.Image.Input( "image_a", tooltip="First source image.", ), io.Image.Input( "image_b", tooltip="Second source image.", ), io.Mask.Input( "mask", tooltip="Blend mask (white = image_a, black = image_b).", ), io.Boolean.Input( "trim_to_shortest", default=True, tooltip="Trim image_a, image_b, and mask to the shortest sequence length before blending.", ), io.Int.Input( "mask_low_res_dilation", default=5, min=0, max=15, tooltip="Downscale the mask to long side 64, dilate it spatially, then resize it back before blending.", ), ], outputs=[ io.Image.Output("image"), ], ) @classmethod def execute( cls, image_a: torch.Tensor, image_b: torch.Tensor, mask: torch.Tensor, trim_to_shortest: bool, mask_low_res_dilation: int, ) -> io.NodeOutput: device = comfy.model_management.get_torch_device() if mask.ndim == 4: mask = mask[:, :, :, 0] if mask.ndim == 2: mask = mask.unsqueeze(0) if image_a.shape[1:3] != image_b.shape[1:3]: raise ValueError( "image_a and image_b must have the same spatial resolution" ) if image_a.shape[1:3] != mask.shape[1:3]: raise ValueError( "image_a, image_b, and mask must have the same spatial resolution" ) if trim_to_shortest: shortest = min(image_a.shape[0], image_b.shape[0], mask.shape[0]) image_a = image_a[:shortest] image_b = image_b[:shortest] mask = mask[:shortest] elif not (image_a.shape[0] == image_b.shape[0] == mask.shape[0]): raise ValueError( "image_a, image_b, and mask must have the same frame count unless trim_to_shortest is enabled" ) original = image_a.permute(0, 3, 1, 2) target = image_b.permute(0, 3, 1, 2) mask_4d = mask.unsqueeze(1).float() mask_4d = _apply_low_res_mask_dilation(mask_4d, mask_low_res_dilation) result = _pyramid_blend( original, target, mask_4d, max_level=7, device=device, output_device=comfy.model_management.intermediate_device(), ) return io.NodeOutput(result.permute(0, 2, 3, 1))