154 lines
5.2 KiB
Python
154 lines
5.2 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
from comfy_api.latest import io
|
|
|
|
from .nodes_registry import comfy_node
|
|
|
|
|
|
@comfy_node(name="LTXVDilateVideoMask")
|
|
class LTXVDilateVideoMask(io.ComfyNode):
|
|
"""Dilates a video mask spatially and/or temporally using max-pooling and thresholds the result."""
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="LTXVDilateVideoMask",
|
|
category="Lightricks/mask_operations",
|
|
description=(
|
|
"Dilates a video mask spatially and/or temporally using "
|
|
"separable max-pooling and thresholds the result."
|
|
),
|
|
inputs=[
|
|
io.Int.Input(
|
|
"spatial_radius",
|
|
default=1,
|
|
min=0,
|
|
max=300,
|
|
tooltip="Half-size of the spatial dilation kernel. Kernel = 2*radius+1.",
|
|
),
|
|
io.Int.Input(
|
|
"temporal_radius",
|
|
default=0,
|
|
min=0,
|
|
max=10,
|
|
tooltip="Half-size of the temporal dilation kernel. Kernel = 2*radius+1.",
|
|
),
|
|
io.Mask.Input(
|
|
"mask",
|
|
optional=True,
|
|
tooltip="Video mask to dilate. Either this or image_as_mask must be provided.",
|
|
),
|
|
io.Image.Input(
|
|
"image_as_mask",
|
|
optional=True,
|
|
tooltip="Image to use as mask (channel-averaged). Either this or mask must be provided.",
|
|
),
|
|
],
|
|
outputs=[
|
|
io.Mask.Output("mask"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(
|
|
cls,
|
|
spatial_radius: int,
|
|
temporal_radius: int,
|
|
mask: torch.Tensor | None = None,
|
|
image_as_mask: torch.Tensor | None = None,
|
|
) -> io.NodeOutput:
|
|
if mask is None and image_as_mask is None:
|
|
raise ValueError("Either 'mask' or 'image_as_mask' must be provided.")
|
|
|
|
if mask is None:
|
|
mask = image_as_mask.mean(dim=-1)
|
|
|
|
if mask.ndim == 4:
|
|
mask = mask[:, :, :, 0]
|
|
|
|
s_kernel = spatial_radius * 2 + 1
|
|
t_kernel = temporal_radius * 2 + 1
|
|
|
|
# Separable dilation: 2D spatial + 1D temporal (much faster than 3D pooling)
|
|
if s_kernel > 1:
|
|
mask = mask.unsqueeze(1) # (B, 1, H, W)
|
|
mask = F.max_pool2d(
|
|
mask, kernel_size=s_kernel, stride=1, padding=spatial_radius
|
|
)
|
|
mask = mask.squeeze(1)
|
|
|
|
if t_kernel > 1:
|
|
B, H, W = mask.shape
|
|
mask = mask.permute(1, 2, 0).reshape(H * W, 1, B)
|
|
mask = F.max_pool1d(
|
|
mask, kernel_size=t_kernel, stride=1, padding=temporal_radius
|
|
)
|
|
mask = mask.reshape(H, W, B).permute(2, 0, 1)
|
|
|
|
mask = (mask > 0.5).float()
|
|
return io.NodeOutput(mask)
|
|
|
|
|
|
_BG_COLOR_RGB = (102, 255, 0)
|
|
|
|
|
|
@comfy_node(name="LTXVInpaintPreprocess")
|
|
class LTXVInpaintPreprocess(io.ComfyNode):
|
|
"""Composites images with a green (#66FF00) background where mask is active.
|
|
|
|
If the mask has a single frame it is broadcast to match the video length.
|
|
|
|
The fill is always SDR display-encoded ``#66FF00`` (``_BG_COLOR_RGB``), even
|
|
when ``images`` are already ACEScct. Native HDR inpaint IC-LoRA was trained
|
|
with that pixel-space key; converting the green into ACEScct changes the
|
|
VAE latent cue and breaks inpainting.
|
|
"""
|
|
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="LTXVInpaintPreprocess",
|
|
category="Lightricks/image_processing",
|
|
description=(
|
|
"Composites images with a green (#66FF00) background where mask "
|
|
"is active, for inpainting conditioning. Always paints SDR "
|
|
"#66FF00 (not ACEScct-converted) — that key matches IC-LoRA "
|
|
"training even on native HDR / ACEScct plates."
|
|
),
|
|
inputs=[
|
|
io.Image.Input(
|
|
"images",
|
|
tooltip="Video frames to composite onto the green background.",
|
|
),
|
|
io.Mask.Input(
|
|
"mask",
|
|
tooltip="Mask indicating regions to replace with green. Single-frame masks are broadcast.",
|
|
),
|
|
],
|
|
outputs=[
|
|
io.Image.Output("image"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(
|
|
cls,
|
|
images: torch.Tensor,
|
|
mask: torch.Tensor,
|
|
) -> io.NodeOutput:
|
|
if mask.ndim == 4:
|
|
mask = mask[:, :, :, 0]
|
|
|
|
if mask.shape[0] == 1 and images.shape[0] > 1:
|
|
mask = mask.expand(images.shape[0], -1, -1)
|
|
|
|
min_frames = min(mask.shape[0], images.shape[0])
|
|
mask = mask[:min_frames]
|
|
images = images[:min_frames]
|
|
|
|
mask_4d = mask.unsqueeze(-1) # (B, H, W, 1) for broadcasting
|
|
bg_color = torch.tensor(_BG_COLOR_RGB).float().to(images.device) / 255
|
|
|
|
result = images * (1 - mask_4d) + bg_color.view(1, 1, 1, 3) * mask_4d
|
|
return io.NodeOutput(result)
|