diff --git a/__init__.py b/__init__.py index 70583c3..bdbe7f3 100644 --- a/__init__.py +++ b/__init__.py @@ -9,6 +9,8 @@ from .vace_batch_context import WanVACEBatchContext from .frame_number_overlay import FrameNumberOverlay from .vace_outpaint import VACEOutpaint from .vace_inpaint import WanVACEInpaint +from .wan_first_last_middle_frame import WanFirstLastMiddleFrameToVideo +from .vace_first_middle_last import WanVACEFirstMiddleLast NODE_CLASS_MAPPINGS = { "WanVACEPrep": WanVACEPrep, @@ -19,6 +21,8 @@ NODE_CLASS_MAPPINGS = { "FrameNumberOverlay": FrameNumberOverlay, "VACEOutpaint": VACEOutpaint, "WanVACEInpaint": WanVACEInpaint, + "WanFirstLastMiddleFrameToVideo": WanFirstLastMiddleFrameToVideo, + "WanVACEFirstMiddleLast": WanVACEFirstMiddleLast, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -30,6 +34,8 @@ NODE_DISPLAY_NAME_MAPPINGS = { "FrameNumberOverlay": "🪐 Frame Number Overlay", "VACEOutpaint": "🪐 Video Outpaint", "WanVACEInpaint": "🪐 Wan VACE Inpaint (Experimental)", + "WanFirstLastMiddleFrameToVideo": "Wan First/Last/Middle Frame to Video", + "WanVACEFirstMiddleLast": "🪐 VACE First/Middle/Last", } WEB_DIRECTORY = "./web" diff --git a/vace_first_middle_last.py b/vace_first_middle_last.py new file mode 100644 index 0000000..e9e8a20 --- /dev/null +++ b/vace_first_middle_last.py @@ -0,0 +1,101 @@ +import torch +import torch.nn.functional as F + + +def _resize_frames(frames, width, height): + if frames.shape[1] == height and frames.shape[2] == width: + return frames + x = frames.movedim(-1, 1).float() + x = F.interpolate(x, size=(height, width), mode="bilinear", align_corners=False) + return x.movedim(1, -1).to(frames.dtype) + + +class WanVACEFirstMiddleLast: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "width": ("INT", {"default": 832, "min": 16, "max": 16384, "step": 16}), + "height": ("INT", {"default": 480, "min": 16, "max": 16384, "step": 16}), + "length": ("INT", { + "default": 81, "min": 1, "max": 16384, "step": 4, + "tooltip": "Total frame count. Must follow the 4n+1 pattern (1, 5, 9, ..., 81, ...)." + }), + "middle_position": ("FLOAT", { + "default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, + "tooltip": "Where middle frames are centered in the video, as a fraction of total length." + }), + }, + "optional": { + "first": ("IMAGE",), + "middle": ("IMAGE",), + "last": ("IMAGE",), + } + } + + RETURN_TYPES = ("IMAGE", "MASK", "INT", "INT", "INT") + RETURN_NAMES = ("control_video", "control_mask", "width", "height", "length") + FUNCTION = "run" + CATEGORY = "video/VACE" + DESCRIPTION = ( + "Builds a VACE control video and mask from optional first, middle, and last " + "frame batches. Known frames are placed at their positions with mask=0; " + "remaining frames are gray placeholders with mask=1 for generation." + ) + + def run(self, width, height, length, middle_position, first=None, middle=None, last=None): + snapped_width = (width // 16) * 16 + snapped_height = (height // 16) * 16 + snapped_length = ((length - 1) // 4) * 4 + 1 + + if snapped_width != width or snapped_height != height: + print( + f"[WanVACEFirstMiddleLast] Dimensions snapped to 16px grid: " + f"{width}x{height} -> {snapped_width}x{snapped_height}" + ) + if snapped_length != length: + print( + f"[WanVACEFirstMiddleLast] length snapped to 4n+1: " + f"{length} -> {snapped_length}" + ) + + width = snapped_width + height = snapped_height + length = snapped_length + + control = torch.full((length, height, width, 3), 0.5, dtype=torch.float32) + mask = torch.ones((length, height, width), dtype=torch.float32) + + n_first = 0 + n_last = 0 + + if first is not None: + first = _resize_frames(first, width, height) + n_first = min(first.shape[0], length) + control[:n_first] = first[:n_first, :, :, :3] + mask[:n_first] = 0.0 + + if last is not None: + last = _resize_frames(last, width, height) + n_last = min(last.shape[0], length - n_first) + if n_last > 0: + control[length - n_last:] = last[-n_last:, :, :, :3] + mask[length - n_last:] = 0.0 + + if middle is not None: + middle = _resize_frames(middle, width, height) + zone_start = n_first + zone_end = length - n_last + available = zone_end - zone_start + + if available > 0: + n_mid = min(middle.shape[0], available) + target_center = round(middle_position * (length - 1)) + mid_start = target_center - n_mid // 2 + mid_start = max(mid_start, zone_start) + mid_start = min(mid_start, zone_end - n_mid) + mid_end = mid_start + n_mid + control[mid_start:mid_end] = middle[:n_mid, :, :, :3] + mask[mid_start:mid_end] = 0.0 + + return (control, mask, width, height, length) diff --git a/wan_first_last_middle_frame.py b/wan_first_last_middle_frame.py new file mode 100644 index 0000000..62e2d07 --- /dev/null +++ b/wan_first_last_middle_frame.py @@ -0,0 +1,126 @@ +import torch +import comfy.model_management +import comfy.utils +import comfy.clip_vision +import node_helpers + + +class WanFirstLastMiddleFrameToVideo: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "positive": ("CONDITIONING",), + "negative": ("CONDITIONING",), + "vae": ("VAE",), + "width": ("INT", {"default": 832, "min": 16, "max": 16384, "step": 16}), + "height": ("INT", {"default": 480, "min": 16, "max": 16384, "step": 16}), + "length": ("INT", {"default": 81, "min": 1, "max": 16384, "step": 4}), + "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}), + }, + "optional": { + "start_image": ("IMAGE",), + "end_image": ("IMAGE",), + "middle_image": ("IMAGE",), + "middle_frame": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "clip_vision_start_image": ("CLIP_VISION_OUTPUT",), + "clip_vision_end_image": ("CLIP_VISION_OUTPUT",), + "clip_vision_middle_image": ("CLIP_VISION_OUTPUT",), + }, + } + + RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT") + RETURN_NAMES = ("positive", "negative", "latent") + FUNCTION = "execute" + CATEGORY = "wan_vace_prep" + + def execute(self, positive, negative, vae, width, height, length, batch_size, + start_image=None, end_image=None, middle_image=None, middle_frame=0.5, + clip_vision_start_image=None, clip_vision_end_image=None, + clip_vision_middle_image=None): + spacial_scale = vae.spacial_compression_encode() + latent = torch.zeros( + [batch_size, vae.latent_channels, ((length - 1) // 4) + 1, + height // spacial_scale, width // spacial_scale], + device=comfy.model_management.intermediate_device() + ) + + if start_image is not None: + start_image = comfy.utils.common_upscale( + start_image[:length].movedim(-1, 1), width, height, "bilinear", "center" + ).movedim(1, -1) + if end_image is not None: + end_image = comfy.utils.common_upscale( + end_image[-length:].movedim(-1, 1), width, height, "bilinear", "center" + ).movedim(1, -1) + if middle_image is not None: + middle_image = comfy.utils.common_upscale( + middle_image[:length].movedim(-1, 1), width, height, "bilinear", "center" + ).movedim(1, -1) + + image = torch.ones((length, height, width, 3)) * 0.5 + mask = torch.ones((1, 1, latent.shape[2] * 4, latent.shape[-2], latent.shape[-1])) + + if start_image is not None: + image[:start_image.shape[0]] = start_image + mask[:, :, :start_image.shape[0] + 3] = 0.0 + + if end_image is not None: + image[-end_image.shape[0]:] = end_image + mask[:, :, -end_image.shape[0]:] = 0.0 + + if middle_image is not None: + mid_start = int(((middle_frame * length) // 4) * 4) + mid_start = min(mid_start, length - 1) + mid_end = min(mid_start + middle_image.shape[0], length) + image[mid_start:mid_end] = middle_image[:mid_end - mid_start] + mask[:, :, mid_start + 3:mid_end + 3] = 0.0 + + concat_latent_image = vae.encode(image[:, :, :, :3]) + mask = mask.view(1, mask.shape[2] // 4, 4, mask.shape[3], mask.shape[4]).transpose(1, 2) + positive = node_helpers.conditioning_set_values( + positive, {"concat_latent_image": concat_latent_image, "concat_mask": mask} + ) + negative = node_helpers.conditioning_set_values( + negative, {"concat_latent_image": concat_latent_image, "concat_mask": mask} + ) + + clip_vision_output = None + if clip_vision_start_image is not None: + clip_vision_output = clip_vision_start_image + + if clip_vision_middle_image is not None: + if clip_vision_output is not None: + states = torch.cat( + [clip_vision_output.penultimate_hidden_states, + clip_vision_middle_image.penultimate_hidden_states], + dim=-2 + ) + clip_vision_output = comfy.clip_vision.Output() + clip_vision_output.penultimate_hidden_states = states + else: + clip_vision_output = clip_vision_middle_image + + if clip_vision_end_image is not None: + if clip_vision_output is not None: + states = torch.cat( + [clip_vision_output.penultimate_hidden_states, + clip_vision_end_image.penultimate_hidden_states], + dim=-2 + ) + clip_vision_output = comfy.clip_vision.Output() + clip_vision_output.penultimate_hidden_states = states + else: + clip_vision_output = clip_vision_end_image + + if clip_vision_output is not None: + positive = node_helpers.conditioning_set_values( + positive, {"clip_vision_output": clip_vision_output} + ) + negative = node_helpers.conditioning_set_values( + negative, {"clip_vision_output": clip_vision_output} + ) + + out_latent = {} + out_latent["samples"] = latent + return (positive, negative, out_latent)