import base64 import json from io import BytesIO import numpy as np import torch from comfy_api.latest import io from PIL import Image from .nodes_registry import comfy_node def _catmull_rom(p0: dict, p1: dict, p2: dict, p3: dict, t: float) -> dict[str, float]: t2 = t * t t3 = t2 * t return { "x": 0.5 * ( 2 * p1["x"] + (-p0["x"] + p2["x"]) * t + (2 * p0["x"] - 5 * p1["x"] + 4 * p2["x"] - p3["x"]) * t2 + (-p0["x"] + 3 * p1["x"] - 3 * p2["x"] + p3["x"]) * t3 ), "y": 0.5 * ( 2 * p1["y"] + (-p0["y"] + p2["y"]) * t + (2 * p0["y"] - 5 * p1["y"] + 4 * p2["y"] - p3["y"]) * t2 + (-p0["y"] + 3 * p1["y"] - 3 * p2["y"] + p3["y"]) * t3 ), } def _interpolate_spline( control_points: list[dict], num_samples: int ) -> list[dict[str, int]]: """Catmull-Rom spline interpolation matching the JS frontend logic.""" if len(control_points) == 0: return [] if len(control_points) == 1: p = control_points[0] return [{"x": round(p["x"]), "y": round(p["y"])} for _ in range(num_samples)] if len(control_points) == 2: a, b = control_points return [ { "x": round(a["x"] + (b["x"] - a["x"]) * i / (num_samples - 1)), "y": round(a["y"] + (b["y"] - a["y"]) * i / (num_samples - 1)), } for i in range(num_samples) ] pts = [control_points[0], *control_points, control_points[-1]] n_seg = len(pts) - 3 result = [] for i in range(num_samples): g_t = (i / (num_samples - 1)) * n_seg seg = min(int(g_t), n_seg - 1) l_t = g_t - seg p = _catmull_rom(pts[seg], pts[seg + 1], pts[seg + 2], pts[seg + 3], l_t) result.append({"x": round(p["x"]), "y": round(p["y"])}) return result @comfy_node(name="LTXVSparseTrackEditor", description="LTX Sparse Track Editor") class LTXVSparseTrackEditor(io.ComfyNode): """Interactive spline editor for drawing sparse motion tracks. Provides a canvas widget where users can draw and edit spline control points on top of a reference image. Outputs interpolated track coordinates compatible with LTXVDrawTracks. """ @classmethod def define_schema(cls): return io.Schema( node_id="LTXVSparseTrackEditor", category="Lightricks/motion_tracking", description=( "Interactive spline editor for drawing sparse motion tracks " "on a reference image." ), inputs=[ io.Image.Input( "image", tooltip="Reference image displayed as the editor canvas background.", ), io.String.Input( "points_store", default="[]", tooltip="JSON array of spline control points managed by the editor widget.", ), io.String.Input( "coordinates", default="[]", tooltip="JSON array of interpolated track coordinates produced by the editor.", ), io.Int.Input( "points_to_sample", default=121, min=2, max=10000, tooltip="Number of points sampled along each spline curve.", ), ], outputs=[ io.String.Output("tracks"), ], is_output_node=True, ) @classmethod def execute( cls, image, points_store: str, coordinates: str, points_to_sample: int, ) -> io.NodeOutput: # Re-interpolate from control points so that changes to # points_to_sample are always respected, regardless of JS sync. try: splines = json.loads(points_store) if points_store else [] except (json.JSONDecodeError, TypeError): splines = [] if splines and isinstance(splines, list) and isinstance(splines[0], list): interpolated = [_interpolate_spline(sp, points_to_sample) for sp in splines] tracks = json.dumps(interpolated) elif coordinates and coordinates != "[]": tracks = coordinates else: tracks = "[]" img_array = (image[0].cpu().numpy() * 255).astype(np.uint8) img = Image.fromarray(img_array) buf = BytesIO() img.save(buf, format="JPEG", quality=75) img_b64 = base64.b64encode(buf.getvalue()).decode("utf-8") return io.NodeOutput(tracks, ui={"bg_image": [img_b64]}) def _parse_tracks(raw: str) -> list[list[dict]]: """Parse tracks from a JSON string, handling nested/wrapped formats.""" parsed = json.loads(raw) if isinstance(raw, str) else raw if isinstance(parsed, list): unwrapped = [] for item in parsed: unwrapped.append(json.loads(item) if isinstance(item, str) else item) parsed = unwrapped tracks: list[list[dict]] = [] stack = [parsed] while stack: obj = stack.pop() if isinstance(obj, list) and len(obj) > 0: if isinstance(obj[0], dict) and "x" in obj[0] and "y" in obj[0]: tracks.append(obj) else: stack.extend(obj) return tracks def _age_color_batch(ratios: torch.Tensor, device: torch.device) -> torch.Tensor: """Vectorised age-ratio -> RGB [0..1] mapping on GPU. Gradient: blue -> green -> yellow -> red. """ colors = torch.zeros(ratios.shape[0], 3, device=device) m1 = ratios <= 1 / 3 tr1 = ratios[m1] * 3 colors[m1, 1] = tr1 colors[m1, 2] = 1 - tr1 m2 = (ratios > 1 / 3) & (ratios <= 2 / 3) tr2 = (ratios[m2] - 1 / 3) * 3 colors[m2, 0] = tr2 colors[m2, 1] = 1 m3 = ratios > 2 / 3 tr3 = (ratios[m3] - 2 / 3) * 3 colors[m3, 0] = 1 colors[m3, 1] = 1 - tr3 return colors def _render_resolution(width: int, height: int, reference_short_side: int): """Compute the higher render resolution that preserves aspect ratio.""" if height <= width: rw = int(width * reference_short_side / height) rh = reference_short_side else: rw = reference_short_side rh = int(height * reference_short_side / width) scale_x = rw / width scale_y = rh / height return rw, rh, scale_x, scale_y _MIN_RADIUS = 2 _MAX_RADIUS = 8 _MAX_TRAIL = 50 _REF_SHORT_SIDE = 1080 @comfy_node(name="LTXVDrawTracks", description="LTX Draw Sparse Tracks") class LTXVDrawTracks(io.ComfyNode): """GPU-accelerated sparse track renderer. Renders circles at a high reference resolution and downscales with bilinear interpolation so circle sizes match the CPU version. All work — rasterisation, compositing and resize — stays on GPU. """ @classmethod def define_schema(cls): return io.Schema( node_id="LTXVDrawTracks", category="Lightricks/motion_tracking", description=( "GPU-accelerated sparse track renderer. Rasterises circles at " "high resolution and downscales with bilinear interpolation." ), inputs=[ io.String.Input( "tracks", multiline=True, tooltip="JSON string of track coordinates (list of point lists with x/y keys).", ), io.Int.Input( "width", default=512, min=8, max=8192, step=8, tooltip="Output image width in pixels.", ), io.Int.Input( "height", default=512, min=8, max=8192, step=8, tooltip="Output image height in pixels.", ), ], outputs=[ io.Image.Output(), ], ) @classmethod def execute(cls, tracks: str, width: int, height: int) -> io.NodeOutput: device = torch.device("cuda" if torch.cuda.is_available() else "cpu") parsed = _parse_tracks(tracks) if not parsed: blank = torch.zeros(1, height, width, 3, device=device, dtype=torch.half) return io.NodeOutput(blank) num_tracks = len(parsed) num_frames = max(len(t) for t in parsed) rw, rh, sx, sy = _render_resolution(width, height, _REF_SHORT_SIDE) point_xy = torch.zeros(num_tracks, num_frames, 2, device=device) vis = torch.zeros(num_tracks, num_frames, dtype=torch.bool, device=device) for i, trk in enumerate(parsed): coords = torch.tensor( [[p["x"] * sx, p["y"] * sy] for p in trk], dtype=torch.float32, device=device, ) point_xy[i, : len(trk)] = coords vis[i, : len(trk)] = True max_d = 2 * _MAX_RADIUS + 3 half_d = max_d // 2 offsets = torch.arange(max_d, device=device) - half_d oy, ox = torch.meshgrid(offsets, offsets, indexing="ij") template_dist_sq = oy.float().square() + ox.float().square() render_frames = torch.zeros(num_frames, rh, rw, 3, device=device) for t in range(num_frames): tau_min = max(0, t - _MAX_TRAIL) window = t - tau_min + 1 active_xy = point_xy[:, tau_min : t + 1] active_vis = vis[:, tau_min : t + 1] ages = torch.arange(window - 1, -1, -1, device=device, dtype=torch.float32) ratios = 1.0 - ages / _MAX_TRAIL radii = _MIN_RADIUS + (_MAX_RADIUS - _MIN_RADIUS) * ratios colors = _age_color_batch(ratios, device) flat_xy = active_xy.reshape(-1, 2) flat_vis = active_vis.reshape(-1) flat_radii = radii.unsqueeze(0).expand(num_tracks, -1).reshape(-1) flat_colors = colors.unsqueeze(0).expand(num_tracks, -1, -1).reshape(-1, 3) idx = flat_vis.nonzero(as_tuple=True)[0] if idx.shape[0] == 0: continue pts = flat_xy[idx] r = flat_radii[idx] c = flat_colors[idx] flat_ages = ages.unsqueeze(0).expand(num_tracks, -1).reshape(-1) sort_order = flat_ages[idx].argsort(descending=True) pts = pts[sort_order] r = r[sort_order] c = c[sort_order] _rasterise_circles( render_frames[t], pts, r, c, template_dist_sq, half_d, max_d, rh, rw ) out = torch.nn.functional.interpolate( render_frames.permute(0, 3, 1, 2), size=(height, width), mode="bilinear", align_corners=False, ).permute(0, 2, 3, 1) out = out[..., [2, 1, 0]] # RGB -> BGR to match IC-LoRA training data format return io.NodeOutput(out.half()) def _rasterise_circles( frame: torch.Tensor, pts: torch.Tensor, radii: torch.Tensor, colors: torch.Tensor, template_dist_sq: torch.Tensor, half_d: int, max_d: int, H: int, W: int, ) -> None: """Stamp filled circles onto *frame* fully on-device. Uses ``scatter_reduce_`` with ``'amax'`` to resolve overlaps in painter's order (circles are expected oldest-first so the highest index = newest wins). """ M = pts.shape[0] if M == 0: return device = pts.device # per-circle masks [M, D, D] radii_sq = (radii * radii).view(M, 1, 1) circle_masks = template_dist_sq.unsqueeze(0) <= radii_sq # frame-space indices [M, D, D] cx = pts[:, 0].round().long().view(M, 1, 1) cy = pts[:, 1].round().long().view(M, 1, 1) offsets_y = torch.arange(max_d, device=device).sub(half_d).view(1, max_d, 1) offsets_x = torch.arange(max_d, device=device).sub(half_d).view(1, 1, max_d) fy = (cy + offsets_y).expand(M, max_d, max_d) # [M, D, D] fx = (cx + offsets_x).expand(M, max_d, max_d) # [M, D, D] valid = circle_masks & (fy >= 0) & (fy < H) & (fx >= 0) & (fx < W) flat_fy = fy[valid] flat_fx = fx[valid] flat_lin = (flat_fy * W + flat_fx).long() # circle index per valid pixel (oldest=0 … newest=M-1) j_map = torch.arange(M, device=device, dtype=torch.float32).view(M, 1, 1) j_map = j_map.expand_as(valid) flat_j = j_map[valid] # priority map — highest index (newest) wins via 'amax' reduce priority = torch.full((H * W,), -1.0, device=device) priority.scatter_reduce_(0, flat_lin, flat_j, reduce="amax", include_self=False) priority = priority.view(H, W).long() has_circle = priority >= 0 frame[has_circle] = colors[priority[has_circle]]