From 2d05008f4adde54ded504bda380735aefa6c3ebc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?St=C3=A9phane=20du=20Hamel?= Date: Sun, 3 Aug 2025 21:47:43 +0200 Subject: [PATCH] Wan trackToVideo tiled --- __init__.py | 2 ++ nodes_wan.py | 95 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 97 insertions(+) diff --git a/__init__.py b/__init__.py index 2d25932..7cfae46 100644 --- a/__init__.py +++ b/__init__.py @@ -8,6 +8,7 @@ NODE_CLASS_MAPPINGS = { "WanVaceToVideoTiledVAE": WanVaceToVideoTiledVAE, "WanCameraImageToVideoTiledVAE":WanCameraImageToVideoTiledVAE, "WanPhantomSubjectToVideoTiledVAE": WanPhantomSubjectToVideoTiledVAE, + "WanTrackToVideoTiledVAE": WanTrackToVideoTiledVAE, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -18,4 +19,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "WanVaceToVideoTiledVAE": "WanVaceToVideo (Tiled VAE encode)", "WanCameraImageToVideoTiledVAE": "WanCameraImageToVideo (Tiled VAE encode)", "WanPhantomSubjectToVideoTiledVAE": "WanPhantomSubjectToVideo (Tiled VAE encode)", + "WanTrackToVideoTiledVAE": "WanTrackToVideo (Tiled VAE encode)", } diff --git a/nodes_wan.py b/nodes_wan.py index 33143f6..f4578d1 100644 --- a/nodes_wan.py +++ b/nodes_wan.py @@ -1,5 +1,6 @@ from comfy.sd import VAE # import math +from comfy_extras.nodes_wan import pad_pts, parse_json_tracks, patch_motion, process_tracks import nodes import node_helpers import torch @@ -400,3 +401,97 @@ class WanPhantomSubjectToVideoTiledVAE: out_latent = {} out_latent["samples"] = latent return (positive, cond2, negative, out_latent) + + +class WanTrackToVideoTiledVAE: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "positive": ("CONDITIONING", ), + "negative": ("CONDITIONING", ), + "vae": ("VAE", ), + "tracks": ("STRING", {"multiline": True, "default": "[]"}), + "width": ("INT", {"default": 832, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}), + "height": ("INT", {"default": 480, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}), + "length": ("INT", {"default": 81, "min": 1, "max": nodes.MAX_RESOLUTION, "step": 4}), + "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}), + "temperature": ("FLOAT", {"default": 220.0, "min": 1.0, "max": 1000.0, "step": 0.1}), + "topk": ("INT", {"default": 2, "min": 1, "max": 10}), + "start_image": ("IMAGE", ), + + "tile_size": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 64}), + "overlap": ("INT", {"default": 64, "min": 0, "max": 4096, "step": 32}), + "temporal_size": ("INT", {"default": 64, "min": 8, "max": 4096, "step": 4, "tooltip": "Amount of frames to encode at a time."}), + "temporal_overlap": ("INT", {"default": 8, "min": 4, "max": 4096, "step": 4, "tooltip": "Amount of frames to overlap."}) + }, + "optional": { + "clip_vision_output": ("CLIP_VISION_OUTPUT", ), + }} + + RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT") + RETURN_NAMES = ("positive", "negative", "latent") + FUNCTION = "encode" + + CATEGORY = "conditioning/video_models" + + def encode(self, positive, negative, vae, tracks, width, height, length, batch_size, + temperature, topk, tile_size, overlap, temporal_size=64, temporal_overlap=8, start_image=None, clip_vision_output=None): + + tracks_data = parse_json_tracks(tracks) + + if not tracks_data: + return WanImageToVideoTiledVAE().encode(positive, negative, vae, width, height, length, batch_size, tile_size, overlap, temporal_size=temporal_size, temporal_overlap=temporal_overlap, start_image=start_image, clip_vision_output=clip_vision_output) + + latent = torch.zeros([batch_size, 16, ((length - 1) // 4) + 1, height // 8, width // 8], + device=comfy.model_management.intermediate_device()) + + if isinstance(tracks_data[0][0], dict): + tracks_data = [tracks_data] + + processed_tracks = [] + for batch in tracks_data: + arrs = [] + for track in batch: + pts = pad_pts(track) + arrs.append(pts) + + tracks_np = np.stack(arrs, axis=0) + processed_tracks.append(process_tracks(tracks_np, (width, height), length - 1).unsqueeze(0)) + + if start_image is not None: + start_image = comfy.utils.common_upscale(start_image[:batch_size].movedim(-1, 1), width, height, "bilinear", "center").movedim(1, -1) + videos = torch.ones((start_image.shape[0], length, height, width, start_image.shape[-1]), device=start_image.device, dtype=start_image.dtype) * 0.5 + for i in range(start_image.shape[0]): + videos[i, 0] = start_image[i] + + latent_videos = [] + videos = comfy.utils.resize_to_batch_size(videos, batch_size) + for i in range(batch_size): + latent_videos += [vae.encode_tiled(videos[i, :, :, :, :3], tile_x=tile_size, tile_y=tile_size, overlap=overlap, tile_t=temporal_size, overlap_t=temporal_overlap)] + y = torch.cat(latent_videos, dim=0) + + # Scale latent since patch_motion is non-linear + y = comfy.latent_formats.Wan21().process_in(y) + + processed_tracks = comfy.utils.resize_list_to_batch_size(processed_tracks, batch_size) + res = patch_motion( + processed_tracks, y, temperature=temperature, topk=topk, vae_divide=(4, 16) # type: ignore + ) + + mask, concat_latent_image = res + concat_latent_image = comfy.latent_formats.Wan21().process_out(concat_latent_image) + mask = -mask + 1.0 # Invert mask to match expected format + positive = node_helpers.conditioning_set_values(positive, + {"concat_mask": mask, + "concat_latent_image": concat_latent_image}) + negative = node_helpers.conditioning_set_values(negative, + {"concat_mask": mask, + "concat_latent_image": concat_latent_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)