Fixed alpha output

This commit is contained in:
malkuthro
2025-10-07 21:26:31 +02:00
parent 8960af66c3
commit a90c399caf
+54 -31
View File
@@ -3,9 +3,11 @@ import torch
class MultiKeyframeImageBatch:
"""
A custom ComfyUI node for VFX workflows that takes 1-4 keyframe images with specified frame positions
and creates an IMAGE batch sequence of a given total length, filling undefined frames with black images.
Useful for preparing sparse control sequences for video models like Wan 2.2, where black frames indicate
no latent update. Assumes all input images have identical dimensions.
and creates an IMAGE batch sequence of a given total length, filling undefined frames with selected placeholder color (black or 50% gray).
Also outputs a MASK batch where occupied frames are white (1.0) and empty are black (0.0), directly compatible with standard mask nodes.
Supports VFX-style frame numbering with a start_frame offset (internal indices are 0-based).
Useful for preparing sparse control sequences for video models like Wan 2.2, where placeholder frames indicate no latent update.
Assumes all input images have identical dimensions.
"""
@classmethod
@@ -19,84 +21,105 @@ class MultiKeyframeImageBatch:
"step": 1,
"display": "number"
}),
"image1": ("IMAGE", ),
"frame1": ("INT", {
"start_frame": ("INT", {
"default": 0,
"min": 0,
"max": 1023,
"max": 999999,
"step": 1,
"display": "number"
}),
"placeholder_color": (["Black", "Gray"], {
"default": "Black"
}),
"image1": ("IMAGE", ),
"image1_frame": ("INT", {
"default": 0,
"min": 0,
"max": 999999,
"step": 1,
"display": "number"
}),
},
"optional": {
"image2": ("IMAGE", ),
"frame2": ("INT", {
"image2_frame": ("INT", {
"default": 4,
"min": 0,
"max": 1023,
"max": 999999,
"step": 1,
"display": "number"
}),
"image3": ("IMAGE", ),
"frame3": ("INT", {
"image3_frame": ("INT", {
"default": 8,
"min": 0,
"max": 1023,
"max": 999999,
"step": 1,
"display": "number"
}),
"image4": ("IMAGE", ),
"frame4": ("INT", {
"image4_frame": ("INT", {
"default": 12,
"min": 0,
"max": 1023,
"max": 999999,
"step": 1,
"display": "number"
}),
}
}
RETURN_TYPES = ("IMAGE", )
RETURN_NAMES = ("image_batch", )
RETURN_TYPES = ("IMAGE", "MASK", )
RETURN_NAMES = ("image_batch", "mask_batch", )
FUNCTION = "create_batch"
CATEGORY = "Koolook/VFX"
OUTPUT_NODE = False # Not necessarily an output node, but can be chained
def create_batch(self, total_frames, image1, frame1, image2=None, frame2=None, image3=None, frame3=None, image4=None, frame4=None):
def create_batch(self, total_frames, start_frame, placeholder_color, image1, image1_frame, image2=None, image2_frame=None, image3=None, image3_frame=None, image4=None, image4_frame=None):
if image1.shape[0] != 1:
raise ValueError("Each input should be a single IMAGE (batch size 1), not a pre-batched tensor.")
h, w, c = image1.shape[1:]
black = torch.zeros(h, w, c) # [H, W, C] zero tensor for black frame
if placeholder_color == "Black":
placeholder = torch.zeros(h, w, c) # [H, W, C] zero tensor for black
else: # Gray
placeholder = torch.full((h, w, c), 0.5) # 50% gray in 0-1 float space
# Initialize list of frames with blacks
frame_images = [black.clone() for _ in range(total_frames)]
# For masks: single-channel [H, W]
black_mask = torch.zeros((h, w), dtype=torch.float32)
white_mask = torch.ones((h, w), dtype=torch.float32)
# Initialize lists for frames and masks
frame_images = [placeholder.clone() for _ in range(total_frames)]
frame_masks = [black_mask.clone() for _ in range(total_frames)]
# Collect connected keyframes
keyframes = [(image1[0], frame1)]
if image2 is not None and frame2 is not None:
keyframes = [(image1[0], image1_frame)]
if image2 is not None and image2_frame is not None:
if image2.shape[1:] != (h, w, c):
raise ValueError("All images must have the same height, width, and channels.")
keyframes.append((image2[0], frame2))
if image3 is not None and frame3 is not None:
keyframes.append((image2[0], image2_frame))
if image3 is not None and image3_frame is not None:
if image3.shape[1:] != (h, w, c):
raise ValueError("All images must have the same height, width, and channels.")
keyframes.append((image3[0], frame3))
if image4 is not None and frame4 is not None:
keyframes.append((image3[0], image3_frame))
if image4 is not None and image4_frame is not None:
if image4.shape[1:] != (h, w, c):
raise ValueError("All images must have the same height, width, and channels.")
keyframes.append((image4[0], frame4))
keyframes.append((image4[0], image4_frame))
# Place keyframes in the sequence
for img, f in keyframes:
if 0 <= f < total_frames:
frame_images[f] = img
index = f - start_frame
if 0 <= index < total_frames:
frame_images[index] = img
frame_masks[index] = white_mask.clone()
else:
print(f"Warning: Frame {f} is out of range (0 to {total_frames-1}); ignoring this keyframe.")
print(f"Warning: Frame {f} maps to index {index} which is out of range (0 to {total_frames-1}); ignoring this keyframe.")
# Stack into batch tensor [total_frames, H, W, C]
batch = torch.stack(frame_images, dim=0)
return (batch, )
# Stack into batch tensors
image_batch = torch.stack(frame_images, dim=0) # [total_frames, H, W, C]
mask_batch = torch.stack(frame_masks, dim=0) # [total_frames, H, W]
return (image_batch, mask_batch, )
# Node mappings for ComfyUI registration
NODE_CLASS_MAPPINGS = {