267 lines
8.3 KiB
Python
267 lines
8.3 KiB
Python
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))
|