Add LTXRefineAlphaMatte for post-process of IC-LoRA alpha mattes.
Cleanup (expand/choke, feather, clip, unmix) plus optional plate composite and a live preview, identity at default widgets. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
obergman
@@ -40,6 +40,7 @@ from .low_vram_loaders import (
|
||||
LowVRAMLatentUpscaleModelLoader,
|
||||
)
|
||||
from .masks import LTXVPreprocessMasks
|
||||
from .matte_refine import LTXRefineAlphaMatte
|
||||
from .nodes_registry import NODE_CLASS_MAPPINGS as RUNTIME_NODE_CLASS_MAPPINGS
|
||||
from .nodes_registry import (
|
||||
NODE_DISPLAY_NAME_MAPPINGS as RUNTIME_NODE_DISPLAY_NAME_MAPPINGS,
|
||||
@@ -120,6 +121,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
"LTXVHDRDecodePostprocess": LTXVHDRDecodePostprocess,
|
||||
"LTXVAudioOnlyModel": LTXVAudioOnlyModel,
|
||||
"LTXVAudioOnlyEmptyVideoLatent": LTXVAudioOnlyEmptyVideoLatent,
|
||||
"LTXRefineAlphaMatte": LTXRefineAlphaMatte,
|
||||
}
|
||||
|
||||
# Consistent display names between static and dynamic node mappings in nodes_registry.py,
|
||||
|
||||
@@ -0,0 +1,202 @@
|
||||
"""Shared matte → RGBA / over-plate helpers for Refine and Over Background."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def nchw_mask(mask: torch.Tensor) -> torch.Tensor:
|
||||
mask = mask.to(dtype=torch.float32)
|
||||
if mask.ndim == 2:
|
||||
mask = mask.unsqueeze(0)
|
||||
return mask
|
||||
|
||||
|
||||
def rgb_frames(image: torch.Tensor) -> torch.Tensor:
|
||||
if image.ndim == 3:
|
||||
image = image.unsqueeze(0)
|
||||
return image[..., :3]
|
||||
|
||||
|
||||
def frames_from_video(video: Any) -> torch.Tensor:
|
||||
if hasattr(video, "get_components"):
|
||||
images = video.get_components().images
|
||||
else:
|
||||
images = video
|
||||
return rgb_frames(images)
|
||||
|
||||
|
||||
def luma_mask(image: torch.Tensor) -> torch.Tensor:
|
||||
"""IMAGE matte → MASK. 4-channel uses alpha; otherwise the red channel
|
||||
(same as Comfy ImageToMask `red`, which matches a grayscale RGB matte)."""
|
||||
if image.ndim == 3:
|
||||
image = image.unsqueeze(0)
|
||||
image = image.to(dtype=torch.float32)
|
||||
if image.shape[-1] >= 4:
|
||||
return image[..., 3].clamp(0.0, 1.0)
|
||||
return image[..., 0].clamp(0.0, 1.0)
|
||||
|
||||
|
||||
def mask_from_matte(
|
||||
matte: torch.Tensor | None,
|
||||
mask: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
if matte is not None:
|
||||
return luma_mask(matte)
|
||||
if mask is None:
|
||||
raise ValueError("Connect matte frames (IMAGE) or a MASK.")
|
||||
return nchw_mask(mask)
|
||||
|
||||
|
||||
def resize_rgb(image: torch.Tensor, height: int, width: int) -> torch.Tensor:
|
||||
if image.shape[1] == height and image.shape[2] == width:
|
||||
return image
|
||||
x = image.movedim(-1, 1)
|
||||
x = F.interpolate(x, size=(height, width), mode="bilinear", align_corners=False)
|
||||
return x.movedim(1, -1)
|
||||
|
||||
|
||||
def resize_mask(mask: torch.Tensor, height: int, width: int) -> torch.Tensor:
|
||||
if mask.shape[-2:] == (height, width):
|
||||
return mask
|
||||
return F.interpolate(
|
||||
mask.unsqueeze(1),
|
||||
size=(height, width),
|
||||
mode="bilinear",
|
||||
align_corners=False,
|
||||
).squeeze(1)
|
||||
|
||||
|
||||
def match_length(plate: torch.Tensor, frames: int, length: str) -> torch.Tensor:
|
||||
count = plate.shape[0]
|
||||
if count == frames:
|
||||
return plate
|
||||
if count > frames:
|
||||
return plate[:frames]
|
||||
if length == "loop" and count > 1:
|
||||
reps = int(math.ceil(frames / count))
|
||||
return plate.repeat(reps, *([1] * (plate.ndim - 1)))[:frames]
|
||||
extra = frames - count
|
||||
return torch.cat([plate, plate[-1:].expand(extra, *([-1] * (plate.ndim - 1)))], dim=0)
|
||||
|
||||
|
||||
def match_mask_to_rgb(
|
||||
alpha: torch.Tensor,
|
||||
frames: int,
|
||||
height: int,
|
||||
width: int,
|
||||
length: str,
|
||||
) -> torch.Tensor:
|
||||
if alpha.shape[-2:] != (height, width):
|
||||
alpha = resize_mask(alpha, height, width)
|
||||
if alpha.shape[0] == 1 and frames > 1:
|
||||
return alpha.expand(frames, -1, -1)
|
||||
if alpha.shape[0] != frames:
|
||||
return match_length(alpha.unsqueeze(-1), frames, length).squeeze(-1)
|
||||
return alpha
|
||||
|
||||
|
||||
def fit_plate(plate: torch.Tensor, height: int, width: int, fit: str) -> torch.Tensor:
|
||||
plate = plate[..., :3]
|
||||
n, h, w, _ = plate.shape
|
||||
if h == height and w == width:
|
||||
return plate
|
||||
if fit == "stretch":
|
||||
return resize_rgb(plate, height, width)
|
||||
|
||||
if fit == "contain":
|
||||
scale = min(height / h, width / w)
|
||||
else:
|
||||
scale = max(height / h, width / w)
|
||||
nh = max(1, int(round(h * scale)))
|
||||
nw = max(1, int(round(w * scale)))
|
||||
if fit == "contain":
|
||||
nh = min(nh, height)
|
||||
nw = min(nw, width)
|
||||
else:
|
||||
nh = max(nh, height)
|
||||
nw = max(nw, width)
|
||||
x = resize_rgb(plate, nh, nw)
|
||||
if fit == "contain":
|
||||
canvas = torch.zeros((n, height, width, 3), device=x.device, dtype=x.dtype)
|
||||
top = (height - nh) // 2
|
||||
left = (width - nw) // 2
|
||||
canvas[:, top : top + nh, left : left + nw] = x
|
||||
return canvas
|
||||
top = max(0, (nh - height) // 2)
|
||||
left = max(0, (nw - width) // 2)
|
||||
cropped = x[:, top : top + height, left : left + width]
|
||||
if cropped.shape[1] != height or cropped.shape[2] != width:
|
||||
return resize_rgb(plate, height, width)
|
||||
return cropped
|
||||
|
||||
|
||||
def pad_even(image: torch.Tensor) -> torch.Tensor:
|
||||
"""H.264 yuv420p cannot encode odd width or height."""
|
||||
n, h, w, c = image.shape
|
||||
ph = h + (h % 2)
|
||||
pw = w + (w % 2)
|
||||
if ph == h and pw == w:
|
||||
return image
|
||||
out = image.new_zeros((n, ph, pw, c))
|
||||
out[:, :h, :w] = image
|
||||
return out
|
||||
|
||||
|
||||
def resolve_plate(
|
||||
background: str,
|
||||
height: int,
|
||||
width: int,
|
||||
frames: int,
|
||||
fit: str,
|
||||
length: str,
|
||||
device: torch.device,
|
||||
dtype: torch.dtype,
|
||||
background_image: torch.Tensor | None = None,
|
||||
background_video: Any | None = None,
|
||||
) -> torch.Tensor | None:
|
||||
if background == "none":
|
||||
return None
|
||||
if background == "image":
|
||||
if background_image is None:
|
||||
raise ValueError(
|
||||
"background=image but no still is loaded. Use Load plate on this node, "
|
||||
"or connect background_image."
|
||||
)
|
||||
plate = rgb_frames(background_image).to(device=device, dtype=dtype)
|
||||
else:
|
||||
if background_video is None:
|
||||
raise ValueError(
|
||||
"background=video but no clip is loaded. Use Load plate on this node, "
|
||||
"or connect background_video."
|
||||
)
|
||||
plate = frames_from_video(background_video).to(device=device, dtype=dtype)
|
||||
plate = fit_plate(plate, height, width, fit)
|
||||
return match_length(plate, frames, length)
|
||||
|
||||
|
||||
def over_composite(
|
||||
fg: torch.Tensor,
|
||||
alpha: torch.Tensor,
|
||||
plate: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
fg = fg[..., :3].contiguous()
|
||||
if plate is None:
|
||||
return pad_even(fg)
|
||||
if fg.shape[0] == 1 and plate.shape[0] > 1:
|
||||
fg = fg.expand(plate.shape[0], -1, -1, -1)
|
||||
alpha = alpha.expand(plate.shape[0], -1, -1)
|
||||
a = alpha.unsqueeze(-1).to(device=fg.device, dtype=fg.dtype)
|
||||
out = fg * a + plate * (1.0 - a)
|
||||
return pad_even(out.clamp(0.0, 1.0).contiguous())
|
||||
|
||||
|
||||
def join_rgba(rgb: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor:
|
||||
"""Straight RGBA. Alpha 1 = keep subject (does not invert like Join Image with Alpha)."""
|
||||
rgb = rgb[..., :3]
|
||||
a = alpha.unsqueeze(-1).to(device=rgb.device, dtype=rgb.dtype).clamp(0.0, 1.0)
|
||||
return torch.cat([rgb, a], dim=-1).contiguous()
|
||||
+514
@@ -0,0 +1,514 @@
|
||||
"""User-facing matte cleanup after IC-LoRA alpha generation."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
|
||||
import folder_paths
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from comfy_api.latest import InputImpl, io
|
||||
from PIL import Image
|
||||
|
||||
from .alpha_composite import (
|
||||
join_rgba,
|
||||
match_length,
|
||||
match_mask_to_rgb,
|
||||
over_composite,
|
||||
mask_from_matte,
|
||||
resolve_plate,
|
||||
resize_rgb,
|
||||
rgb_frames,
|
||||
)
|
||||
from .nodes_registry import NODES_DISPLAY_NAME_PREFIX, comfy_node
|
||||
|
||||
_LIVE_MAX_SIDE = 320
|
||||
_LIVE_MAX_FRAMES = 240
|
||||
_PLATE_CACHE: dict[tuple[str, int], torch.Tensor] = {}
|
||||
|
||||
|
||||
def _live_b64(frame: torch.Tensor, jpeg: bool = False) -> str:
|
||||
t = frame.detach().float().cpu()
|
||||
if t.ndim == 2:
|
||||
img = Image.fromarray((t.clamp(0, 1) * 255).byte().numpy(), mode="L")
|
||||
fmt, kw = "PNG", {"compress_level": 6}
|
||||
else:
|
||||
img = Image.fromarray((t[..., :3].clamp(0, 1) * 255).byte().numpy(), mode="RGB")
|
||||
fmt, kw = ("JPEG", {"quality": 60}) if jpeg else ("PNG", {"compress_level": 6})
|
||||
buf = BytesIO()
|
||||
img.save(buf, format=fmt, **kw)
|
||||
return base64.b64encode(buf.getvalue()).decode("ascii")
|
||||
|
||||
|
||||
def _index_time(tensor: torch.Tensor, idx: torch.Tensor) -> torch.Tensor:
|
||||
return tensor.index_select(0, idx.to(device=tensor.device))
|
||||
|
||||
|
||||
def _live_clip(
|
||||
source: torch.Tensor,
|
||||
rgb: torch.Tensor,
|
||||
plate: torch.Tensor | None = None,
|
||||
) -> tuple[list[str], list[str], list[str], float, int]:
|
||||
n_source = int(source.shape[0])
|
||||
n, h, w = n_source, int(source.shape[-2]), int(source.shape[-1])
|
||||
if n > _LIVE_MAX_FRAMES:
|
||||
idx = torch.linspace(0, n - 1, _LIVE_MAX_FRAMES, device=source.device).round().long()
|
||||
source = _index_time(source, idx)
|
||||
rgb = _index_time(rgb, idx)
|
||||
if plate is not None:
|
||||
plate = _index_time(plate, idx)
|
||||
n = int(source.shape[0])
|
||||
scale = min(1.0, _LIVE_MAX_SIDE / float(max(h, w)))
|
||||
if scale < 1.0:
|
||||
nh = max(1, int(round(h * scale)))
|
||||
nw = max(1, int(round(w * scale)))
|
||||
source = F.interpolate(
|
||||
source.unsqueeze(1), size=(nh, nw), mode="bilinear", align_corners=False
|
||||
).squeeze(1)
|
||||
rgb = resize_rgb(rgb, nh, nw)
|
||||
if plate is not None:
|
||||
plate = resize_rgb(plate, nh, nw)
|
||||
source_cpu = source.detach().float().cpu()
|
||||
rgb_cpu = rgb.detach().float().cpu()
|
||||
masks = [_live_b64(source_cpu[i], jpeg=False) for i in range(n)]
|
||||
rgbs = [_live_b64(rgb_cpu[i], jpeg=True) for i in range(n)]
|
||||
bgs: list[str] = []
|
||||
if plate is not None:
|
||||
plate_cpu = plate.detach().float().cpu()
|
||||
bgs = [_live_b64(plate_cpu[i], jpeg=True) for i in range(n)]
|
||||
return masks, rgbs, bgs, float(scale), n_source
|
||||
|
||||
|
||||
def _expand_mask(mask: torch.Tensor, pixels: int) -> torch.Tensor:
|
||||
"""Square dilate (positive) or erode (negative) in pixel units.
|
||||
|
||||
Replicate-pad before the pool so the border matches the live preview (which
|
||||
clamps coordinates) instead of max_pool2d's zero fill.
|
||||
"""
|
||||
if pixels == 0:
|
||||
return mask
|
||||
x = mask.unsqueeze(1)
|
||||
pad = abs(int(pixels))
|
||||
k = pad * 2 + 1
|
||||
if pixels < 0:
|
||||
x = 1.0 - x
|
||||
x = F.pad(x, (pad, pad, 0, 0), mode="replicate")
|
||||
x = F.max_pool2d(x, kernel_size=(1, k), stride=1)
|
||||
x = F.pad(x, (0, 0, pad, pad), mode="replicate")
|
||||
x = F.max_pool2d(x, kernel_size=(k, 1), stride=1)
|
||||
if pixels < 0:
|
||||
x = 1.0 - x
|
||||
return x.squeeze(1)
|
||||
|
||||
|
||||
def _gaussian_blur(mask: torch.Tensor, sigma: float) -> torch.Tensor:
|
||||
if sigma <= 0:
|
||||
return mask
|
||||
radius = max(1, int(round(float(sigma) * 3.0)))
|
||||
k = radius * 2 + 1
|
||||
coords = torch.arange(k, device=mask.device, dtype=mask.dtype) - radius
|
||||
g = torch.exp(-(coords * coords) / (2.0 * sigma * sigma))
|
||||
g = g / g.sum()
|
||||
x = mask.unsqueeze(1)
|
||||
x = F.pad(x, (0, 0, radius, radius), mode="replicate")
|
||||
x = F.conv2d(x, g.view(1, 1, k, 1))
|
||||
x = F.pad(x, (radius, radius, 0, 0), mode="replicate")
|
||||
x = F.conv2d(x, g.view(1, 1, 1, k))
|
||||
return x.squeeze(1)
|
||||
|
||||
|
||||
def _gaussian_blur_rgb(image: torch.Tensor, sigma: float) -> torch.Tensor:
|
||||
"""NHWC RGB, same separable blur as `_gaussian_blur`."""
|
||||
if sigma <= 0:
|
||||
return image
|
||||
n, h, w, c = image.shape
|
||||
x = image.movedim(-1, 1).reshape(n * c, h, w)
|
||||
return _gaussian_blur(x, sigma).reshape(n, c, h, w).movedim(1, -1)
|
||||
|
||||
|
||||
_UNMIX_FLOOR = 0.1
|
||||
_UNMIX_GAIN_FLOOR = 0.25
|
||||
|
||||
|
||||
def _uniform_clamp(rgb: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
|
||||
"""Move `rgb` toward `target` as far as stays in gamut, scaling all channels
|
||||
together. Per-channel clamping would clip one channel first and shift hue."""
|
||||
delta = target - rgb
|
||||
big = torch.full_like(delta, 1e6)
|
||||
to_high = torch.where(delta > 0, (1.0 - rgb) / delta.clamp_min(1e-6), big)
|
||||
to_low = torch.where(delta < 0, rgb / (-delta).clamp_min(1e-6), big)
|
||||
t = torch.minimum(to_high, to_low).amin(dim=-1, keepdim=True).clamp(0.0, 1.0)
|
||||
return rgb + t * delta
|
||||
|
||||
|
||||
def _unmix_color(
|
||||
rgb: torch.Tensor,
|
||||
alpha: torch.Tensor,
|
||||
amount: float,
|
||||
radius: float,
|
||||
) -> torch.Tensor:
|
||||
"""Recover straight foreground color: C = αF + (1-α)B.
|
||||
|
||||
Estimates B from low-alpha pixels, then F = (C - (1-α)B) / α, which is already
|
||||
identity at α=1. Two guards keep a bad α from turning into a colored fringe:
|
||||
the divisor is floored at `_UNMIX_GAIN_FLOOR` so background removal is never
|
||||
amplified more than ~3x, and the result is pulled back into gamut uniformly
|
||||
instead of per channel. Below `_UNMIX_FLOOR` the correction fades out, since
|
||||
α is then too small to solve for F at all. Amount 0 is identity.
|
||||
"""
|
||||
if amount <= 0:
|
||||
return rgb
|
||||
a = alpha.clamp(0.0, 1.0)
|
||||
bg_w = (1.0 - a) * (1.0 - a)
|
||||
sigma = max(0.5, float(radius))
|
||||
num = _gaussian_blur_rgb(rgb * bg_w.unsqueeze(-1), sigma)
|
||||
den = _gaussian_blur(bg_w, sigma).unsqueeze(-1).clamp_min(1e-4)
|
||||
background = (num / den).clamp(0.0, 1.0)
|
||||
safe = a.unsqueeze(-1).clamp_min(_UNMIX_GAIN_FLOOR)
|
||||
foreground = _uniform_clamp(rgb, (rgb - (1.0 - safe) * background) / safe)
|
||||
gate = (a / _UNMIX_FLOOR).clamp(0.0, 1.0).unsqueeze(-1)
|
||||
mix = float(amount) * gate
|
||||
return (rgb * (1.0 - mix) + foreground * mix).clamp(0.0, 1.0)
|
||||
|
||||
|
||||
def _clip_levels(mask: torch.Tensor, black_clip: float, white_clip: float) -> torch.Tensor:
|
||||
lo = float(black_clip)
|
||||
hi = 1.0 - float(white_clip)
|
||||
if lo <= 0.0 and hi >= 1.0:
|
||||
return mask
|
||||
if hi <= lo:
|
||||
hi = lo + 1e-4
|
||||
return ((mask - lo) / (hi - lo)).clamp(0.0, 1.0)
|
||||
|
||||
|
||||
def _refine_mask(
|
||||
mask: torch.Tensor,
|
||||
invert: bool,
|
||||
expand: int,
|
||||
feather: float,
|
||||
black_clip: float,
|
||||
white_clip: float,
|
||||
mix: float,
|
||||
) -> torch.Tensor:
|
||||
work = mask
|
||||
if invert:
|
||||
work = 1.0 - work
|
||||
raw = work
|
||||
mix = min(1.0, max(0.0, float(mix)))
|
||||
if mix <= 0.0:
|
||||
return raw.clamp(0.0, 1.0)
|
||||
work = _expand_mask(work, int(expand))
|
||||
work = _clip_levels(work, black_clip, white_clip)
|
||||
work = _gaussian_blur(work, float(feather))
|
||||
if mix < 1.0:
|
||||
work = torch.lerp(raw, work, mix)
|
||||
return work.clamp(0.0, 1.0)
|
||||
|
||||
|
||||
def _load_plate_file(name: str) -> torch.Tensor:
|
||||
"""Read a still or video from Comfy's input folder as NHWC RGB."""
|
||||
name = (name or "").strip().replace("\\", "/")
|
||||
parts = Path(name).parts
|
||||
if not name or Path(name).is_absolute() or ".." in parts:
|
||||
raise ValueError(f"Invalid plate file: {name!r}")
|
||||
exists = getattr(folder_paths, "exists_annotated_filepath", None)
|
||||
if callable(exists) and not exists(name):
|
||||
raise ValueError(f"Plate file not found: {name}")
|
||||
path = folder_paths.get_annotated_filepath(name)
|
||||
try:
|
||||
mtime = Path(path).stat().st_mtime_ns
|
||||
except OSError:
|
||||
mtime = 0
|
||||
key = (str(path), int(mtime))
|
||||
cached = _PLATE_CACHE.get(key)
|
||||
if cached is not None:
|
||||
return cached
|
||||
images = InputImpl.VideoFromFile(path).get_components().images
|
||||
if images is None or int(images.shape[0]) == 0:
|
||||
raise ValueError(f"Could not read plate file: {name}")
|
||||
frames = rgb_frames(images)
|
||||
_PLATE_CACHE.clear()
|
||||
_PLATE_CACHE[key] = frames
|
||||
return frames
|
||||
|
||||
|
||||
@comfy_node(name="LTXRefineAlphaMatte")
|
||||
class LTXRefineAlphaMatte(io.ComfyNode):
|
||||
"""Post-process a generated alpha matte. Identity at default widget values."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="LTXRefineAlphaMatte",
|
||||
display_name=NODES_DISPLAY_NAME_PREFIX + " Refine Alpha Matte",
|
||||
category="Lightricks/IC-LoRA",
|
||||
description=(
|
||||
"Cleanup for IC-LoRA alpha mattes, plus optional over-plate composite. "
|
||||
"Defaults leave the matte and RGB unchanged. Connect Decode via Get Video "
|
||||
"Components to `matte` (white = subject); `mask` still works. Expand is "
|
||||
"Shift Edge (negative = choke). Feather blurs the silhouette. Black/white "
|
||||
"clip is Resolve Clean Black / Clean White. Unmix pulls old-background "
|
||||
"color out of mixed edges. `rgba` is straight alpha. `over` is "
|
||||
"the opaque composite when background is image/video. Use **Load plate** "
|
||||
"on this node (no extra Load Image / Load Video nodes). Live preview uses "
|
||||
"the same plate as Queue Prompt."
|
||||
),
|
||||
inputs=[
|
||||
io.Image.Input(
|
||||
"matte",
|
||||
optional=True,
|
||||
tooltip="Matte frames (white = subject). Connect Get Video Components from Decode. Optional if `mask` is connected.",
|
||||
),
|
||||
io.Image.Input(
|
||||
"image",
|
||||
optional=True,
|
||||
tooltip="Source RGB for unmix, RGBA, and over. Gray placeholder if unconnected.",
|
||||
),
|
||||
io.Boolean.Input(
|
||||
"invert",
|
||||
default=False,
|
||||
label_on="inverted",
|
||||
label_off="normal",
|
||||
tooltip="Swap subject and background.",
|
||||
),
|
||||
io.Int.Input(
|
||||
"expand",
|
||||
display_name="expand (px)",
|
||||
default=0,
|
||||
min=-128,
|
||||
max=128,
|
||||
tooltip="Shift Edge in pixels. Positive grows the subject, negative chokes it (kills halos).",
|
||||
),
|
||||
io.Float.Input(
|
||||
"feather",
|
||||
default=0.0,
|
||||
min=0.0,
|
||||
max=64.0,
|
||||
step=0.5,
|
||||
tooltip="Gaussian blur of the silhouette in pixels. 2–8 is a typical AE feather.",
|
||||
),
|
||||
io.Float.Input(
|
||||
"black_clip",
|
||||
display_name="clean black",
|
||||
default=0.0,
|
||||
min=0.0,
|
||||
max=0.5,
|
||||
step=0.01,
|
||||
tooltip="Crush near-transparent gray specks in the background to 0.",
|
||||
),
|
||||
io.Float.Input(
|
||||
"white_clip",
|
||||
display_name="clean white",
|
||||
default=0.0,
|
||||
min=0.0,
|
||||
max=0.5,
|
||||
step=0.01,
|
||||
tooltip="Fill near-opaque gray holes in the subject to 1.",
|
||||
),
|
||||
io.Float.Input(
|
||||
"mix",
|
||||
default=1.0,
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
step=0.01,
|
||||
tooltip="0 = raw matte (after invert), 1 = fully refined.",
|
||||
),
|
||||
io.Float.Input(
|
||||
"unmix",
|
||||
default=0.0,
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
step=0.01,
|
||||
tooltip=(
|
||||
"Decontaminate edge color. 0 = original RGB. 1 = fully unmix mixed "
|
||||
"pixels so they keep the subject color, not the old background."
|
||||
),
|
||||
),
|
||||
io.Float.Input(
|
||||
"unmix_radius",
|
||||
display_name="unmix radius",
|
||||
default=12.0,
|
||||
min=1.0,
|
||||
max=64.0,
|
||||
step=0.5,
|
||||
tooltip=(
|
||||
"How far (pixels) to look for background color when unmixing. "
|
||||
"Raise for long hair/fur; lower if nearby objects tint the edge."
|
||||
),
|
||||
),
|
||||
io.Combo.Input(
|
||||
"background",
|
||||
options=["none", "image", "video"],
|
||||
default="none",
|
||||
tooltip=(
|
||||
"none = original RGB on `over`. image / video = the file from "
|
||||
"**Load plate** (or a wired background_image / background_video)."
|
||||
),
|
||||
),
|
||||
io.Combo.Input(
|
||||
"fit",
|
||||
options=["cover", "contain", "stretch"],
|
||||
default="cover",
|
||||
tooltip="How the plate fills the frame: cover crops, contain letterboxes, stretch distorts.",
|
||||
),
|
||||
io.Combo.Input(
|
||||
"length",
|
||||
options=["hold", "loop"],
|
||||
default="hold",
|
||||
tooltip="If the plate is shorter than the clip: hold last frame or loop.",
|
||||
),
|
||||
io.String.Input(
|
||||
"plate_file",
|
||||
default="",
|
||||
socketless=True,
|
||||
tooltip="Set by Load plate. Filename in Comfy's input folder.",
|
||||
),
|
||||
io.Mask.Input(
|
||||
"mask",
|
||||
optional=True,
|
||||
advanced=True,
|
||||
tooltip="Optional MASK. Used only if you do not connect `matte`.",
|
||||
),
|
||||
io.Image.Input(
|
||||
"background_image",
|
||||
optional=True,
|
||||
lazy=True,
|
||||
advanced=True,
|
||||
tooltip="Optional still from another node. Load plate does not need this.",
|
||||
),
|
||||
io.Video.Input(
|
||||
"background_video",
|
||||
optional=True,
|
||||
lazy=True,
|
||||
advanced=True,
|
||||
tooltip="Optional clip from another node. Load plate does not need this.",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
io.Image.Output(
|
||||
"over",
|
||||
tooltip="Opaque composite over the plate (or original RGB when background=none). Width/height padded to even for H.264; rgba/mask/image keep the source size.",
|
||||
),
|
||||
io.Image.Output(
|
||||
"rgba",
|
||||
tooltip="Straight RGBA from the refined matte (alpha 1 = keep). Use for WebM / EXR.",
|
||||
),
|
||||
io.Mask.Output("mask", tooltip="Refined matte. White = keep."),
|
||||
io.Image.Output(
|
||||
"image",
|
||||
tooltip="Source RGB after color unmix. Same as input when unmix is 0.",
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def check_lazy_status(
|
||||
cls,
|
||||
matte=None,
|
||||
image=None,
|
||||
invert=False,
|
||||
expand=0,
|
||||
feather=0.0,
|
||||
black_clip=0.0,
|
||||
white_clip=0.0,
|
||||
mix=1.0,
|
||||
unmix=0.0,
|
||||
unmix_radius=12.0,
|
||||
background="none",
|
||||
fit="cover",
|
||||
length="hold",
|
||||
plate_file="",
|
||||
mask=None,
|
||||
background_image=None,
|
||||
background_video=None,
|
||||
**_kwargs,
|
||||
):
|
||||
needed = []
|
||||
if background == "image" and background_image is None and not plate_file:
|
||||
needed.append("background_image")
|
||||
if background == "video" and background_video is None and not plate_file:
|
||||
needed.append("background_video")
|
||||
return needed
|
||||
|
||||
@classmethod
|
||||
def execute(
|
||||
cls,
|
||||
matte: torch.Tensor | None = None,
|
||||
image: torch.Tensor | None = None,
|
||||
invert: bool = False,
|
||||
expand: int = 0,
|
||||
feather: float = 0.0,
|
||||
black_clip: float = 0.0,
|
||||
white_clip: float = 0.0,
|
||||
mix: float = 1.0,
|
||||
unmix: float = 0.0,
|
||||
unmix_radius: float = 12.0,
|
||||
background: str = "none",
|
||||
fit: str = "cover",
|
||||
length: str = "hold",
|
||||
plate_file: str = "",
|
||||
mask: torch.Tensor | None = None,
|
||||
background_image: torch.Tensor | None = None,
|
||||
background_video=None,
|
||||
**_kwargs,
|
||||
) -> io.NodeOutput:
|
||||
source = mask_from_matte(matte, mask)
|
||||
work = _refine_mask(
|
||||
source,
|
||||
invert=invert,
|
||||
expand=int(expand),
|
||||
feather=float(feather),
|
||||
black_clip=float(black_clip),
|
||||
white_clip=float(white_clip),
|
||||
mix=float(mix),
|
||||
)
|
||||
|
||||
n, h, w = work.shape
|
||||
if image is None:
|
||||
rgb = torch.full((n, h, w, 3), 0.5, device=work.device, dtype=work.dtype)
|
||||
else:
|
||||
rgb = resize_rgb(rgb_frames(image).to(device=work.device, dtype=work.dtype), h, w)
|
||||
if rgb.shape[0] != n:
|
||||
rgb = match_length(rgb, n, length)
|
||||
|
||||
work = match_mask_to_rgb(work, n, h, w, length)
|
||||
source = match_mask_to_rgb(source, n, h, w, length)
|
||||
unmixed = _unmix_color(rgb, work, float(unmix), float(unmix_radius))
|
||||
if background != "none" and plate_file:
|
||||
loaded = None
|
||||
if background == "image" and background_image is None:
|
||||
loaded = _load_plate_file(plate_file)
|
||||
background_image = loaded
|
||||
elif background == "video" and background_video is None:
|
||||
loaded = _load_plate_file(plate_file)
|
||||
background_video = loaded
|
||||
plate = resolve_plate(
|
||||
background,
|
||||
h,
|
||||
w,
|
||||
n,
|
||||
fit,
|
||||
length,
|
||||
unmixed.device,
|
||||
unmixed.dtype,
|
||||
background_image=background_image,
|
||||
background_video=background_video,
|
||||
)
|
||||
rgba = join_rgba(unmixed, work)
|
||||
over = over_composite(unmixed, work, plate)
|
||||
live_mask, live_rgb, live_bg, live_scale, live_source = _live_clip(source, rgb, plate)
|
||||
return io.NodeOutput(
|
||||
over,
|
||||
rgba,
|
||||
work,
|
||||
unmixed.contiguous(),
|
||||
ui={
|
||||
"live_mask": live_mask,
|
||||
"live_rgb": live_rgb,
|
||||
"live_bg": live_bg,
|
||||
"live_scale": [live_scale],
|
||||
"live_source_frames": [live_source],
|
||||
},
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user