From 7c72ca796ffeb6ebb1873374dfd17bcdafe22fa1 Mon Sep 17 00:00:00 2001 From: siraxe Date: Fri, 17 Oct 2025 12:28:47 +0100 Subject: [PATCH] ATI coordinates --- .gitignore | 3 +- ATI/motion.py | 42 +++ ATI/motion_patch.py | 142 ++++++++ ATI/nodes.py | 333 ++++++++++++++++++ __init__.py | 5 +- nodes/draw_shapes.py | 167 +++++++-- nodes/power_spline_editor.py | 2 +- web/power_spline_editor/node_size_manager.js | 6 +- .../power_spline_editor.js | 2 +- 9 files changed, 673 insertions(+), 29 deletions(-) create mode 100644 ATI/motion.py create mode 100644 ATI/motion_patch.py create mode 100644 ATI/nodes.py diff --git a/.gitignore b/.gitignore index afc867b..4bbd1b9 100644 --- a/.gitignore +++ b/.gitignore @@ -6,4 +6,5 @@ nodes/__pycache__/ .env 22 wanvideo_cached_samples.pt ref_image.jpg -QWEN.md \ No newline at end of file +QWEN.md +ComfyUI-WanVideoWrapper/ \ No newline at end of file diff --git a/ATI/motion.py b/ATI/motion.py new file mode 100644 index 0000000..3d615e6 --- /dev/null +++ b/ATI/motion.py @@ -0,0 +1,42 @@ +# Copyright (c) 2024-2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Dict, List, Optional, Tuple, Union +import numpy as np +import torch + +def process_tracks(tracks_np: np.ndarray, frame_size: Tuple[int, int], quant_multi: int = 8, **kwargs): + # tracks: shape [t, h, w, 3] => samples align with 24 fps, model trained with 16 fps. + # frame_size: tuple (W, H) + + tracks = torch.from_numpy(tracks_np).float() + + if tracks.shape[1] == 121: + tracks = torch.permute(tracks, (1, 0, 2, 3)) + + tracks, visibles = tracks[..., :2], tracks[..., 2:3] + short_edge = min(*frame_size) + + tracks = tracks - torch.tensor([*frame_size]).type_as(tracks) / 2 + tracks = tracks / short_edge * 2 + + visibles = visibles * 2 - 1 + + trange = torch.linspace(-1, 1, tracks.shape[0]).view(-1, 1, 1, 1).expand(*visibles.shape) + + out_ = torch.cat([trange, tracks, visibles], dim=-1).view(121, -1, 4) + out_0 = out_[:1] + out_l = out_[1:] # 121 => 120 | 1 + out_l = torch.repeat_interleave(out_l, 2, dim=0)[1::3] # 120 => 240 => 80 + return torch.cat([out_0, out_l], dim=0) diff --git a/ATI/motion_patch.py b/ATI/motion_patch.py new file mode 100644 index 0000000..1adf1e5 --- /dev/null +++ b/ATI/motion_patch.py @@ -0,0 +1,142 @@ +# Copyright (c) 2024-2025 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import List, Optional, Tuple, Union +import torch + + +# Refer to https://github.com/Angtian/VoGE/blob/main/VoGE/Utils.py +def ind_sel(target: torch.Tensor, ind: torch.Tensor, dim: int = 1): + """ + :param target: [... (can be k or 1), n > M, ...] + :param ind: [... (k), M] + :param dim: dim to apply index on + :return: sel_target [... (k), M, ...] + """ + assert ( + len(ind.shape) > dim + ), "Index must have the target dim, but get dim: %d, ind shape: %s" % (dim, str(ind.shape)) + + target = target.expand( + *tuple( + [ind.shape[k] if target.shape[k] == 1 else -1 for k in range(dim)] + + [ + -1, + ] + * (len(target.shape) - dim) + ) + ) + + ind_pad = ind + + if len(target.shape) > dim + 1: + for _ in range(len(target.shape) - (dim + 1)): + ind_pad = ind_pad.unsqueeze(-1) + ind_pad = ind_pad.expand(*(-1,) * (dim + 1), *target.shape[(dim + 1) : :]) + + return torch.gather(target, dim=dim, index=ind_pad) + + +def merge_final(vert_attr: torch.Tensor, weight: torch.Tensor, vert_assign: torch.Tensor): + """ + + :param vert_attr: [n, d] or [b, n, d] color or feature of each vertex + :param weight: [b(optional), w, h, M] weight of selected vertices + :param vert_assign: [b(optional), w, h, M] selective index + :return: + """ + target_dim = len(vert_assign.shape) - 1 + if len(vert_attr.shape) == 2: + assert vert_attr.shape[0] > vert_assign.max() + # [n, d] ind: [b(optional), w, h, M]-> [b(optional), w, h, M, d] + # sel_attr = ind_sel( + # vert_attr[(None,) * target_dim], vert_assign.type(torch.long), dim=target_dim + # ) + new_shape = [1] * target_dim + list(vert_attr.shape) + tensor = vert_attr.reshape(new_shape) + sel_attr = ind_sel(tensor, vert_assign.type(torch.long), dim=target_dim) + else: + assert vert_attr.shape[1] > vert_assign.max() + #sel_attr = ind_sel( + # vert_attr[:, *(None,) * (target_dim - 1)], vert_assign.type(torch.long), dim=target_dim + #) + new_shape = [vert_attr.shape[0]] + [1] * (target_dim - 1) + list(vert_attr.shape[1:]) + tensor = vert_attr.reshape(new_shape) + sel_attr = ind_sel(tensor, vert_assign.type(torch.long), dim=target_dim) + + # [b(optional), w, h, M] + final_attr = torch.sum(sel_attr * weight.unsqueeze(-1), dim=-2) + return final_attr + + +def patch_motion( + tracks: torch.FloatTensor, # (B, T, N, 4) + vid: torch.FloatTensor, # (C, T, H, W) + temperature: float = 220.0, + vae_divide: tuple = (4, 16), + topk: int = 2, +): + with torch.no_grad(): + _, T, H, W = vid.shape + N = tracks.shape[2] + _, tracks, visible = torch.split( + tracks, [1, 2, 1], dim=-1 + ) # (B, T, N, 2) | (B, T, N, 1) + tracks_n = tracks / torch.tensor([W / min(H, W), H / min(H, W)], device=tracks.device) + tracks_n = tracks_n.clamp(-1, 1) + visible = visible.clamp(0, 1) + + xx = torch.linspace(-W / min(H, W), W / min(H, W), W) + yy = torch.linspace(-H / min(H, W), H / min(H, W), H) + + grid = torch.stack(torch.meshgrid(yy, xx, indexing="ij")[::-1], dim=-1).to( + tracks.device + ) + + tracks_pad = tracks[:, 1:] + visible_pad = visible[:, 1:] + + visible_align = visible_pad.view(T - 1, 4, *visible_pad.shape[2:]).sum(1) + tracks_align = (tracks_pad * visible_pad).view(T - 1, 4, *tracks_pad.shape[2:]).sum( + 1 + ) / (visible_align + 1e-5) + dist_ = ( + (tracks_align[:, None, None] - grid[None, :, :, None]).pow(2).sum(-1) + ) # T, H, W, N + weight = torch.exp(-dist_ * temperature) * visible_align.clamp(0, 1).view( + T - 1, 1, 1, N + ) + vert_weight, vert_index = torch.topk( + weight, k=min(topk, weight.shape[-1]), dim=-1 + ) + + grid_mode = "bilinear" + point_feature = torch.nn.functional.grid_sample( + vid[vae_divide[0]:].permute(1, 0, 2, 3)[:1], + tracks_n[:, :1].type(vid.dtype), + mode=grid_mode, + padding_mode="zeros", + align_corners=False, + ) + point_feature = point_feature.squeeze(0).squeeze(1).permute(1, 0) # N, C=16 + + out_feature = merge_final(point_feature, vert_weight, vert_index).permute(3, 0, 1, 2) # T - 1, H, W, C => C, T - 1, H, W + out_weight = vert_weight.sum(-1) # T - 1, H, W + + # out feature -> already soft weighted + mix_feature = out_feature + vid[vae_divide[0]:, 1:] * (1 - out_weight.clamp(0, 1)) + + out_feature_full = torch.cat([vid[vae_divide[0]:, :1], mix_feature], dim=1) # C, T, H, W + out_mask_full = torch.cat([torch.ones_like(out_weight[:1]), out_weight], dim=0) # T, H, W + return torch.cat([out_mask_full[None].expand(vae_divide[0], -1, -1, -1), out_feature_full], dim=0) diff --git a/ATI/nodes.py b/ATI/nodes.py new file mode 100644 index 0000000..1d9bd8b --- /dev/null +++ b/ATI/nodes.py @@ -0,0 +1,333 @@ +import json +from .motion import process_tracks +import numpy as np +from typing import List, Tuple +import torch +FIXED_LENGTH = 121 +def pad_pts(tr): + """Convert list of {x,y,v} to (FIXED_LENGTH,1,3) array, padding/truncating. + Respects visibility flags from the input if present.""" + pts = np.array([[p['x'], p['y'], p.get('v', 1)] for p in tr], dtype=np.float32) + n = pts.shape[0] + if n < FIXED_LENGTH: + # Create padding with the last point but with visibility = 0 + last_point = pts[-1].copy() + last_point[2] = 0 # Set visibility to 0 for padding + pad = np.array([last_point] * (FIXED_LENGTH - n), dtype=np.float32) + pts = np.vstack((pts, pad)) + else: + pts = pts[:FIXED_LENGTH] + return pts.reshape(FIXED_LENGTH, 1, 3) + +def age_to_bgr(ratio: float) -> Tuple[int,int,int]: + """ + Map ratio∈[0,1] through: 0→blue, 1/3→green, 2/3→yellow, 1→red. + Returns (B,G,R) for OpenCV. + """ + if ratio <= 1/3: + # blue→green + t = ratio / (1/3) + b = int(255 * (1 - t)) + g = int(255 * t) + r = 0 + elif ratio <= 2/3: + # green→yellow + t = (ratio - 1/3) / (1/3) + b = 0 + g = 255 + r = int(255 * t) + else: + # yellow→red + t = (ratio - 2/3) / (1/3) + b = 0 + g = int(255 * (1 - t)) + r = 255 + return (r, g, b) + +def paint_point_track( + frames: np.ndarray, + point_tracks: np.ndarray, + visibles: np.ndarray, + min_radius: int = 1, + max_radius: int = 6, + max_retain: int = 50 +) -> np.ndarray: + """ + Draws every past point of each track on each frame, with radius and color + interpolated by the point's age (old→small to new→large). + + Args: + frames: [F, H, W, 3] uint8 RGB + point_tracks:[N, F, 2] float32 – (x,y) in pixel coords + visibles: [N, F] bool – visibility mask + min_radius: radius for the very first point (oldest) + max_radius: radius for the current point (newest) + + Returns: + video: [F, H, W, 3] uint8 RGB + """ + import cv2 + num_points, num_frames = point_tracks.shape[:2] + H, W = frames.shape[1:3] + + video = frames.copy() + + for t in range(num_frames): + # start from the original frame + frame = video[t].copy() + + for i in range(num_points): + # draw every past step τ = 0..t + for τ in range(t + 1): + if not visibles[i, τ]: + continue + + if t - τ > max_retain: + continue + + # sub-pixel offset + clamp + x, y = point_tracks[i, τ] + 0.5 + xi = int(np.clip(x, 0, W - 1)) + yi = int(np.clip(y, 0, H - 1)) + + # age‐ratio in [0,1] + if num_frames > 1: + ratio = 1 - float(t - τ) / max_retain + else: + ratio = 1.0 + + # interpolated radius + radius = int(round(min_radius + (max_radius - min_radius) * ratio)) + + # OpenCV draws in BGR order: + color_rgb = age_to_bgr(ratio) + + # filled circle + cv2.circle(frame, (xi, yi), radius, color_rgb, thickness=-1) + + video[t] = frame + + return video + +def parse_json_tracks(tracks): + tracks_data = [] + try: + # If tracks is a string, try to parse it as JSON + if isinstance(tracks, str): + parsed = json.loads(tracks.replace("'", '"')) + tracks_data.extend(parsed) + else: + # If tracks is a list of strings, parse each one + for track_str in tracks: + parsed = json.loads(track_str.replace("'", '"')) + tracks_data.append(parsed) + + # Check if we have a single track (dict with x,y) or a list of tracks + if tracks_data and isinstance(tracks_data[0], dict) and 'x' in tracks_data[0]: + # Single track detected, wrap it in a list + tracks_data = [tracks_data] + elif tracks_data and isinstance(tracks_data[0], list) and tracks_data[0] and isinstance(tracks_data[0][0], dict) and 'x' in tracks_data[0][0]: + # Already a list of tracks, nothing to do + pass + else: + # Unexpected format + print(f"Warning: Unexpected track format: {type(tracks_data[0])}") + + except json.JSONDecodeError as e: + print(f"Error parsing tracks JSON: {e}") + tracks_data = [] + + return tracks_data + +class WanVideoATITracks_v2: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("WANVIDEOMODEL", ), + "tracks": ("STRING",), + "width": ("INT", {"default": 832, "min": 64, "max": 2048, "step": 8, "tooltip": "Width of the image to encode"}), + "height": ("INT", {"default": 480, "min": 64, "max": 29048, "step": 8, "tooltip": "Height of the image to encode"}), + "temperature": ("FLOAT", {"default": 220.0, "min": 0.0, "max": 1000.0, "step": 0.1}), + "topk": ("INT", {"default": 2, "min": 1, "max": 10, "step": 1}), + "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Start percent of the steps to apply ATI"}), + "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "End percent of the steps to apply ATI"}), + }, + } + + RETURN_TYPES = ("WANVIDEOMODEL",) + RETURN_NAMES = ("model",) + FUNCTION = "patchmodel" + CATEGORY = "WanVideoWrapper" + + def patchmodel(self, model, tracks, width, height, temperature, topk, start_percent, end_percent): + tracks_data = parse_json_tracks(tracks) + arrs = [] + for track in tracks_data: + pts = pad_pts(track) + arrs.append(pts) + + tracks_np = np.stack(arrs, axis=0) + + processed_tracks = process_tracks(tracks_np, (width, height)) + + patcher = model.clone() + patcher.model_options["transformer_options"]["ati_tracks"] = processed_tracks.unsqueeze(0) + patcher.model_options["transformer_options"]["ati_temperature"] = temperature + patcher.model_options["transformer_options"]["ati_topk"] = topk + patcher.model_options["transformer_options"]["ati_start_percent"] = start_percent + patcher.model_options["transformer_options"]["ati_end_percent"] = end_percent + + return (patcher,) + +class WanVideoATITracksVisualize_v2: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "images": ("IMAGE",), + "tracks": ("STRING",), + "min_radius": ("INT", {"default": 1, "min": 0, "max": 100, "step": 1, "tooltip": "radius for the very first point (oldest)"}), + "max_radius": ("INT", {"default": 6, "min": 0, "max": 100, "step": 1, "tooltip": "radius for the current point (newest)"}), + "max_retain": ("INT", {"default": 50, "min": 0, "max": 100, "step": 1, "tooltip": "Maximum number of points to retain"}), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("images",) + FUNCTION = "patchmodel" + CATEGORY = "WanVideoWrapper" + + def patchmodel(self, images, tracks, min_radius, max_radius, max_retain): + tracks_data = parse_json_tracks(tracks) + arrs = [] + for track in tracks_data: + pts = pad_pts(track) + arrs.append(pts) + + tracks_np = np.stack(arrs, axis=0) + track = np.repeat(tracks_np, 2, axis=1)[:, ::3] + points = track[:, :, 0, :2].astype(np.float32) + visibles = track[:, :, 0, 2].astype(np.float32) + + if images.shape[0] < points.shape[1]: + repeat_count = (points.shape[1] + images.shape[0] - 1) // images.shape[0] + images = images.repeat(repeat_count, 1, 1, 1) + images = images[:points.shape[1]] + elif images.shape[0] > points.shape[1]: + images = images[:points.shape[1]] + + video_viz = paint_point_track(images.cpu().numpy(), points, visibles, min_radius, max_radius, max_retain) + video_viz = torch.from_numpy(video_viz).float() + + return (video_viz,) + +from comfy import utils +import types +from .motion_patch import patch_motion + +class WanConcatCondPatch: + def __init__(self, tracks, temperature, topk): + self.tracks = tracks + self.temperature = temperature + self.topk = topk + + def __get__(self, obj, objtype=None): + # Create bound method with stored parameters + def wrapped_concat_cond(self_module, *args, **kwargs): + return modified_concat_cond(self_module, self.tracks, self.temperature, self.topk, *args, **kwargs) + return types.MethodType(wrapped_concat_cond, obj) + +def modified_concat_cond(self, tracks, temperature, topk, **kwargs): + noise = kwargs.get("noise", None) + extra_channels = self.diffusion_model.patch_embedding.weight.shape[1] - noise.shape[1] + if extra_channels == 0: + return None + + image = kwargs.get("concat_latent_image", None) + device = kwargs["device"] + + if image is None: + shape_image = list(noise.shape) + shape_image[1] = extra_channels + image = torch.zeros(shape_image, dtype=noise.dtype, layout=noise.layout, device=noise.device) + else: + image = utils.common_upscale(image.to(device), noise.shape[-1], noise.shape[-2], "bilinear", "center") + for i in range(0, image.shape[1], 16): + image[:, i: i + 16] = self.process_latent_in(image[:, i: i + 16]) + image = utils.resize_to_batch_size(image, noise.shape[0]) + + if not self.image_to_video or extra_channels == image.shape[1]: + return image + + if image.shape[1] > (extra_channels - 4): + image = image[:, :(extra_channels - 4)] + + mask = kwargs.get("concat_mask", kwargs.get("denoise_mask", None)) + if mask is None: + mask = torch.zeros_like(noise)[:, :4] + else: + if mask.shape[1] != 4: + mask = torch.mean(mask, dim=1, keepdim=True) + mask = 1.0 - mask + mask = utils.common_upscale(mask.to(device), noise.shape[-1], noise.shape[-2], "bilinear", "center") + if mask.shape[-3] < noise.shape[-3]: + mask = torch.nn.functional.pad(mask, (0, 0, 0, 0, 0, noise.shape[-3] - mask.shape[-3]), mode='constant', value=0) + if mask.shape[1] == 1: + mask = mask.repeat(1, 4, 1, 1, 1) + mask = utils.resize_to_batch_size(mask, noise.shape[0]) + + image_cond = torch.cat((mask, image), dim=1) + image_cond_ati = patch_motion(tracks.to(image_cond.device, image_cond.dtype), image_cond[0], + temperature=temperature, topk=topk) + + return image_cond_ati.unsqueeze(0) + +class WanVideoATI_comfy_v2: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ("MODEL", ), + "width": ("INT", {"default": 832, "min": 64, "max": 2048, "step": 8, "tooltip": "Width of the image to encode"}), + "height": ("INT", {"default": 480, "min": 64, "max": 29048, "step": 8, "tooltip": "Height of the image to encode"}), + "tracks": ("STRING",), + "temperature": ("FLOAT", {"default": 220.0, "min": 0.0, "max": 1000.0, "step": 0.1}), + "topk": ("INT", {"default": 2, "min": 1, "max": 10, "step": 1}), + }, + } + + RETURN_TYPES = ("MODEL",) + RETURN_NAMES = ("model", ) + FUNCTION = "patchcond" + CATEGORY = "WanVideoWrapper" + + def patchcond(self, model, tracks, width, height, temperature, topk): + + tracks_data = parse_json_tracks(tracks) + arrs = [] + for track in tracks_data: + pts = pad_pts(track) + arrs.append(pts) + + tracks_np = np.stack(arrs, axis=0) + + processed_tracks = process_tracks(tracks_np, (width, height)) + + model_clone = model.clone() + model_clone.add_object_patch( + "concat_cond", + WanConcatCondPatch( + processed_tracks.unsqueeze(0), temperature, topk + ).__get__(model.model, model.model.__class__) + ) + + return (model_clone,) + +NODE_CLASS_MAPPINGS = { + "WanVideoATITracks_v2": WanVideoATITracks_v2, + "WanVideoATITracksVisualize_v2": WanVideoATITracksVisualize_v2, + "WanVideoATI_comfy_v2": WanVideoATI_comfy_v2, + } +NODE_DISPLAY_NAME_MAPPINGS = { + "WanVideoATITracks_v2": "WanVideo ATI Tracks v2", + "WanVideoATITracksVisualize_v2": "WanVideo ATI Tracks Visualize v2", + "WanVideoATI_comfy_v2": "WanVideo ATI Comfy v2", + } diff --git a/__init__.py b/__init__.py index 2c863a5..cc104d9 100644 --- a/__init__.py +++ b/__init__.py @@ -7,6 +7,7 @@ from .nodes.nodes import * from .nodes.power_loras import * from .nodes.power_spline_editor import * +from .ATI.nodes import * from .nodes.draw_shapes import * from .nodes.draw_joints import * @@ -37,8 +38,8 @@ current_module = sys.modules[__name__] for name, obj in inspect.getmembers(current_module): if inspect.isclass(obj): module_name = getattr(obj, '__module__', '') - # Check if the class is defined in one of the modules within the .nodes package - if module_name.startswith(__name__ + '.nodes'): + # Check if the class is defined in one of the modules within the .nodes package or ATI package + if module_name.startswith(__name__ + '.nodes') or module_name.startswith(__name__ + '.ATI'): # --- ADDED CHECK: Exclude classes with "Math" in the name --- if "Math" in obj.__name__: continue # Skip this class diff --git a/nodes/draw_shapes.py b/nodes/draw_shapes.py index 3d39893..0719dc0 100644 --- a/nodes/draw_shapes.py +++ b/nodes/draw_shapes.py @@ -59,7 +59,9 @@ Locations are center locations. Allows coordinates outside the frame for 'fly-in "trailing": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), "border_width": ("INT", {"default": 0, "min": 0, "max": 100, "step": 1}), "border_color": ("STRING", {"default": 'black'}), - "frames": ("INT", {"forceInput": True}), + "frames": ("INT", {"forceInput": True},{"default": 121 }), + "ati_fade_start": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), + "ati_p_fade_start": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), "preview_enabled": ("BOOLEAN", {"default": True}), } } @@ -880,7 +882,7 @@ Locations are center locations. Allows coordinates outside the frame for 'fly-in # Main Node Method # ---------------------------- def drawshapemask(self, coordinates, bg_image, - shape_width, shape_height, shape_color, bg_color, blur_radius, shape, intensity, + shape_width, shape_height, shape_color, bg_color, blur_radius, shape, intensity, ati_p_fade_start, ati_fade_start, trailing=1.0, border_width=0, border_color='black', frames=None, preview_enabled=True): """ Main entry point. Orchestrates: @@ -1048,11 +1050,8 @@ Locations are center locations. Allows coordinates outside the frame for 'fly-in # Special cases: no animated coords but static points exist -> set batch size accordingly if not processed_coords_list: if static_points: - if p_coords_use_driver and p_driver_path_processed: - batch_size = total_frames - else: - batch_size = 1 - total_frames = 1 + # Always use total_frames for static points to ensure proper spline generation + batch_size = total_frames processed_coords_list = [] path_pause_frames = [] else: @@ -1175,23 +1174,149 @@ Locations are center locations. Allows coordinates outside the frame for 'fly-in # Return minimal 1x1 pixel preview for efficiency when preview is disabled preview_output = torch.zeros([batch_size, 1, 1, 3], dtype=torch.float32) - # ----- Build output coordinates JSON (basic: first path across frames) ----- + # ----- Build output coordinates JSON (format expected by ATI nodes) ----- + # Follow the same format as WanVideoATITracksVisualize + # Pad to 121 frames, then subsample to 81 frames with visibility in the third component output_coords_json = "[]" + all_coords = [] + FIXED_LENGTH = 121 # Match the FIXED_LENGTH in nodes_g.py + + # Process animated paths (affected by ati_fade_start) if processed_coords_list: try: - first_path_coords = [] - path_start_p, path_end_p = path_pause_frames[0] - for i in range(total_frames): - if i < path_start_p: - coord_index = 0 - elif i >= total_frames - path_end_p: - coord_index = len(processed_coords_list[0]) - 1 - else: - coord_index = i - path_start_p - coord_index = max(0, min(coord_index, len(processed_coords_list[0]) - 1)) - first_path_coords.append(processed_coords_list[0][coord_index]) - output_coords_json = json.dumps(first_path_coords) + # Calculate the fade start frame by multiplying percentage by total frames + fade_start_frame = int(ati_fade_start * total_frames) if ati_fade_start > 0 else 0 + + for path_idx, path_coords in enumerate(processed_coords_list): + path_start_p, path_end_p = path_pause_frames[path_idx] + single_path_coords = [] + for i in range(total_frames): + if i < path_start_p: + coord_index = 0 + elif i >= total_frames - path_end_p: + coord_index = len(path_coords) - 1 + else: + coord_index = i - path_start_p + coord_index = max(0, min(coord_index, len(path_coords) - 1)) + + # Extract x, y and determine visibility + coord = path_coords[coord_index] + + # Determine visibility based on ati_fade_start + if ati_fade_start == 0 or i < fade_start_frame: + visibility = 1 # Visible + else: + visibility = 0 # Invisible after fade_start_frame + + single_path_coords.append({ + "x": int(coord["x"]), + "y": int(coord["y"]), + "v": visibility + }) + + # Pad to FIXED_LENGTH (121 frames) if needed + if len(single_path_coords) < FIXED_LENGTH: + # Pad with the last point repeated + last_point = single_path_coords[-1].copy() + last_point["v"] = 0 # Make padded points invisible + while len(single_path_coords) < FIXED_LENGTH: + single_path_coords.append(last_point.copy()) + elif len(single_path_coords) > FIXED_LENGTH: + # Truncate to FIXED_LENGTH + single_path_coords = single_path_coords[:FIXED_LENGTH] + + # Convert to the format expected by ATI: [x, y, visibility] + # This will be further processed by the ATI node like WanVideoATITracksVisualize + all_coords.append(single_path_coords) except Exception: - output_coords_json = "[]" + pass + + # Process static points (p_coordinates) - affected by ati_p_fade_start + if static_points: + try: + # Calculate the fade start frame by multiplying percentage by total frames + fade_start_frame = int(ati_p_fade_start * total_frames) if ati_p_fade_start > 0 else 0 + + for point in static_points: + single_point_spline = [] + for i in range(total_frames): + if ati_p_fade_start == 0 or i < fade_start_frame: + # Keep the point at its original position + visibility = 1 # Visible + else: + # Keep the point at its last position but set visibility to 0 + visibility = 0 # Invisible + + single_point_spline.append({ + "x": int(point["x"]), + "y": int(point["y"]), + "v": visibility + }) + + # Pad to FIXED_LENGTH (121 frames) if needed + if len(single_point_spline) < FIXED_LENGTH: + # Pad with the last point repeated but invisible + last_point = single_point_spline[-1].copy() + last_point["v"] = 0 # Make padded points invisible + while len(single_point_spline) < FIXED_LENGTH: + single_point_spline.append(last_point.copy()) + elif len(single_point_spline) > FIXED_LENGTH: + # Truncate to FIXED_LENGTH + single_point_spline = single_point_spline[:FIXED_LENGTH] + + all_coords.append(single_point_spline) + except Exception as e: + print(f"Error processing static points: {e}") + pass + + # Ensure we have at least one track to avoid the ATI error + if not all_coords: + # Create a default track at origin if no coordinates exist + default_track = [] + for i in range(FIXED_LENGTH): + default_track.append({ + "x": 0, + "y": 0, + "v": 1 if i < total_frames else 0 # Visible for original frames, invisible for padding + }) + all_coords.append(default_track) + + # Format output as a JSON string that ATI nodes can parse + # Follow the same format as WanVideoATITracksVisualize + try: + # Convert tracks to the format that matches what WanVideoATITracksVisualize expects + # The tracks will be padded to 121 frames and subsampled to 81 frames by the ATI node + clean_tracks = [] + for track in all_coords: + clean_track = [] + for point in track: + # Include visibility if it exists, default to 1 if not + v = int(point.get("v", 1)) + clean_track.append({ + "x": int(point["x"]), + "y": int(point["y"]), + "v": v + }) + clean_tracks.append(clean_track) + + # Format output according to the number of tracks: + # - If there's only one track, output it as a single list: [{...}, {...}] + # - If there are multiple tracks, output as a list of lists: [[{...}], [{...}]] + if len(clean_tracks) == 1: + output_coords_json = json.dumps(clean_tracks[0]) + else: + output_coords_json = json.dumps(clean_tracks) + + # Verify the JSON can be parsed correctly + test_parse = json.loads(output_coords_json) + + # Debug output + print(f"DrawShapeOnPath: Generated {len(clean_tracks)} tracks for ATI nodes (padded to {FIXED_LENGTH} frames)") + for i, track in enumerate(clean_tracks): + print(f" Track {i}: {len(track)} points") + except Exception as e: + # Fallback to empty array if there's an issue + print(f"DrawShapeOnPath: Error generating tracks: {e}") + output_coords_json = "[]" return (out_images, out_masks, output_coords_json, preview_output) diff --git a/nodes/power_spline_editor.py b/nodes/power_spline_editor.py index 0147df4..9700cbd 100644 --- a/nodes/power_spline_editor.py +++ b/nodes/power_spline_editor.py @@ -14,7 +14,7 @@ class PowerSplineEditor: "required": { "points_store": ("STRING", {"default": "[{\"x\":200,\"y\":240},{\"x\":304,\"y\":240}]", "multiline": False}), "coordinates": ("STRING", {"multiline": False}), - "mask_width": ("INT", {"default": 600, "min": 8, "max": 4096, "step": 8}), + "mask_width": ("INT", {"default": 640, "min": 8, "max": 4096, "step": 8}), "mask_height": ("INT", {"default": 480, "min": 8, "max": 4096, "step": 8}), "bg_img": (["None", "A", "B", "C"], {"default": "None", "tooltip": "Select background image to overlay at 60% opacity"}), }, diff --git a/web/power_spline_editor/node_size_manager.js b/web/power_spline_editor/node_size_manager.js index 021bef6..0fb0eb1 100644 --- a/web/power_spline_editor/node_size_manager.js +++ b/web/power_spline_editor/node_size_manager.js @@ -26,16 +26,16 @@ export class NodeSizeManager { buttonBarHeight: options.buttonBarHeight || (LiteGraph.NODE_WIDGET_HEIGHT || 30), // === CANVAS DIMENSIONS === - canvasWidth: options.canvasWidth || 600, + canvasWidth: options.canvasWidth || 640, canvasHeight: options.canvasHeight || 480, canvasExtraWidth: options.canvasExtraWidth || 45, - minCanvasWidth: options.minCanvasWidth || 600, + minCanvasWidth: options.minCanvasWidth || 640, maxCanvasWidth: options.maxCanvasWidth || 8192, minCanvasHeight: options.minCanvasHeight || 480, maxCanvasHeight: options.maxCanvasHeight || 8192, // === NODE CONSTRAINTS === - minNodeWidth: options.minNodeWidth || 600, + minNodeWidth: options.minNodeWidth || 640, minNodeHeight: options.minNodeHeight || 480, }; diff --git a/web/power_spline_editor/power_spline_editor.js b/web/power_spline_editor/power_spline_editor.js index 07bdd8c..faeb12d 100644 --- a/web/power_spline_editor/power_spline_editor.js +++ b/web/power_spline_editor/power_spline_editor.js @@ -1719,7 +1719,7 @@ app.registerExtension({ this.imgData = imgData; try { const size = JSON.stringify(this.imgData).length; - if (size < 600 * 480) { // 1MB limit + if (size < 640 * 480) { // 1MB limit sessionStorage.setItem(`spline-editor-img-${this.uuid}`, JSON.stringify(this.imgData)); } else { console.warn("Spline Editor: Image not saved to session storage because it is too large.", size);