diff --git a/k_easy_batch.py b/k_easy_batch.py index 6084dc7..c506248 100644 --- a/k_easy_batch.py +++ b/k_easy_batch.py @@ -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 = {