diff --git a/node_wrappers/realtimenodes/controls/__init__.py b/node_wrappers/realtimenodes/controls/__init__.py index 6e7633b..3e4bcdd 100644 --- a/node_wrappers/realtimenodes/controls/__init__.py +++ b/node_wrappers/realtimenodes/controls/__init__.py @@ -1 +1,2 @@ # This file intentionally left blank + \ No newline at end of file diff --git a/node_wrappers/realtimenodes/controls/flow_warpers.py b/node_wrappers/realtimenodes/controls/flow_warpers.py new file mode 100644 index 0000000..f8716e9 --- /dev/null +++ b/node_wrappers/realtimenodes/controls/flow_warpers.py @@ -0,0 +1,420 @@ +""" +Optical flow nodes for generating and applying flow-based warping. + +These nodes work with optical flow to create more advanced warping effects. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +import math + +from ....src.realtimenodes.control_base import ControlNodeBase +from ....src.utils.general import AlwaysEqualProxy +from ....src.utils.image import gaussian_blur_2d, flow_to_rgb +from ....src.utils.realtime_flownets import RealTimeFlowNet + + +class OpticalFlowNode(ControlNodeBase): + """ + Generate optical flow between images. + + Computes the motion flow field between two images. + """ + + @classmethod + def INPUT_TYPES(cls): + inputs = super().INPUT_TYPES() + inputs["required"].update({ + "previous_image": ("IMAGE", {"tooltip": "Previous frame image"}), + "current_image": ("IMAGE", {"tooltip": "Current frame image"}), + "flow_scale": ( + "FLOAT", + {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1, + "tooltip": "Scale factor to apply to flow"} + ), + "smoothness": ( + "FLOAT", + {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1, + "tooltip": "Gaussian blur sigma to smooth flow field"} + ), + "visualize": ( + "BOOLEAN", + {"default": False, "tooltip": "Output flow visualization as RGB image"} + ), + }) + return inputs + + RETURN_TYPES = ("FLOW_FIELD", "IMAGE") + RETURN_NAMES = ("flow", "flow_visualization") + FUNCTION = "update" + CATEGORY = "Realtime Nodes/control/flow" + + def __init__(self): + super().__init__() + self.flow_net = None + + def update( + self, + previous_image, + current_image, + flow_scale=1.0, + smoothness=1.0, + visualize=False, + always_execute=True, + unique_id=None + ): + """ + Generate optical flow between two images. + + Args: + previous_image: Previous frame tensor [B, H, W, C] + current_image: Current frame tensor [B, H, W, C] + flow_scale: Factor to scale the flow magnitude + smoothness: Gaussian blur sigma to apply to the flow field + visualize: Whether to output a visualization of the flow field + always_execute: Whether to always execute (from ControlNodeBase) + unique_id: Unique ID for state management (from ControlNodeBase) + + Returns: + Flow field tensor [B, 2, H, W] and optional visualization + """ + # Initialize flow network on first use + if self.flow_net is None: + self.flow_net = RealTimeFlowNet() + self.flow_net.eval() + + # Ensure inputs are on the same device + device = previous_image.device + self.flow_net = self.flow_net.to(device) + + # Convert to grayscale and ensure proper format for flow estimation + # Input is BHWC, convert to BCHW for CNN + prev_gray = torch.mean(previous_image, dim=-1, keepdim=True).permute(0, 3, 1, 2) + curr_gray = torch.mean(current_image, dim=-1, keepdim=True).permute(0, 3, 1, 2) + + # Stack frames for flow estimation + stacked = torch.cat([prev_gray, curr_gray], dim=1) + + # Estimate flow + with torch.no_grad(): + flow = self.flow_net(stacked) + + # Scale the flow + flow = flow * flow_scale + + # Apply Gaussian blur if smoothness > 0 + if smoothness > 0: + flow = gaussian_blur_2d(flow, kernel_size=9, sigma=smoothness) + + # Create flow visualization if requested + if visualize: + flow_rgb = flow_to_rgb(flow) + # Convert back to BHWC for ComfyUI + flow_vis = flow_rgb.permute(0, 2, 3, 1) + else: + # Create empty tensor for visualization + flow_vis = torch.zeros_like(previous_image) + + # Print flow statistics for debugging + with torch.no_grad(): + flow_abs_mean = torch.abs(flow).mean().item() + flow_min = flow.min().item() + flow_max = flow.max().item() + print(f"Flow stats - mean abs: {flow_abs_mean:.4f}, min: {flow_min:.4f}, max: {flow_max:.4f}") + + return (flow, flow_vis) + + +class FlowWarpNode(ControlNodeBase): + """ + Apply optical flow field to warp images or latents. + + Uses flow fields for advanced warping effects. + """ + + @classmethod + def INPUT_TYPES(cls): + inputs = super().INPUT_TYPES() + inputs["required"].update({ + "data_type": ( + ["IMAGE", "LATENT"], + {"default": "IMAGE", "tooltip": "Type of data to warp"} + ), + "data": ( + AlwaysEqualProxy("*"), + {"tooltip": "Image or latent to warp with flow field"} + ), + "flow_field": ( + "FLOW_FIELD", + {"tooltip": "Flow field tensor [B, 2, H, W]"} + ), + "strength": ( + "FLOAT", + {"default": 5.0, "min": 0.0, "max": 50.0, "step": 0.1, + "tooltip": "Strength of the flow warping effect (try higher values if no effect is visible)"} + ), + "mode": ( + ["bilinear", "nearest", "bicubic"], + {"default": "bilinear", "tooltip": "Interpolation mode for sampling"} + ), + "padding_mode": ( + ["zeros", "border", "reflection"], + {"default": "border", "tooltip": "How to handle pixels outside boundaries"} + ), + "align_corners": ( + "BOOLEAN", + {"default": True, "tooltip": "Whether to align corners in sampling grid"} + ), + "debug_mode": ( + ["none", "flow", "grid", "test_pattern"], + {"default": "none", "tooltip": "Visualization mode for debugging"} + ), + }) + return inputs + + RETURN_TYPES = (AlwaysEqualProxy("*"), "IMAGE") + RETURN_NAMES = ("warped_data", "debug_visualization") + FUNCTION = "update" + CATEGORY = "Realtime Nodes/control/flow" + + def create_test_pattern(self, height, width, device): + """Create a test pattern grid image to visualize warping effects""" + # Create checkerboard pattern + y, x = torch.meshgrid(torch.linspace(0, height-1, height, device=device), + torch.linspace(0, width-1, width, device=device), + indexing="ij") + + # Create grid lines + grid_size = 32 + grid_x = (x % grid_size < 2).float() + grid_y = (y % grid_size < 2).float() + grid = grid_x + grid_y + + # Create circles + centers = [(height//4, width//4), (height//4, 3*width//4), + (3*height//4, width//4), (3*height//4, 3*width//4)] + circles = torch.zeros_like(grid) + for cy, cx in centers: + dist = torch.sqrt((y - cy)**2 + (x - cx)**2) + circles += (dist < 30).float() + + # Combine patterns + r = grid + g = (x / width) + b = (y / height) + + # Add circles + r = torch.clamp(r + circles, 0, 1) + + # Stack channels + pattern = torch.stack([r, g, b], dim=-1) + + # Add batch dimension + pattern = pattern.unsqueeze(0) + + return pattern + + def update( + self, + data_type, + data, + flow_field, + strength=5.0, + mode="bilinear", + padding_mode="border", + align_corners=True, + debug_mode="none", + always_execute=True, + unique_id=None + ): + """ + Warp image or latent using optical flow field. + + Args: + data_type: Type of data to warp ('IMAGE' or 'LATENT') + data: Image tensor [B, H, W, C] or latent dict + flow_field: Flow field tensor [B, 2, H, W] + strength: Strength of the flow warping effect + mode: Interpolation mode + padding_mode: Padding mode + align_corners: Whether to align corners in sampling grid + debug_mode: Visualization mode for debugging + always_execute: Whether to always execute (from ControlNodeBase) + unique_id: Unique ID for state management (from ControlNodeBase) + + Returns: + Warped image or latent, and debug visualization + """ + # Print debug info + print(f"FlowWarpNode - data_type: {data_type}, strength: {strength}") + print(f"Flow field shape: {flow_field.shape}, dtype: {flow_field.dtype}, device: {flow_field.device}") + if data_type == "IMAGE": + print(f"Image shape: {data.shape}, dtype: {data.dtype}, device: {data.device}") + + # Flow statistics + with torch.no_grad(): + flow_abs_mean = torch.abs(flow_field).mean().item() + flow_min = flow_field.min().item() + flow_max = flow_field.max().item() + print(f"Flow stats - mean abs: {flow_abs_mean:.4f}, min: {flow_min:.4f}, max: {flow_max:.4f}") + + # Scale flow by strength + scaled_flow = flow_field * strength + + # Create debug visualization based on mode + if data_type == "IMAGE": + if debug_mode == "flow": + # Flow RGB visualization + flow_rgb = flow_to_rgb(scaled_flow) + debug_vis = flow_rgb.permute(0, 2, 3, 1) # BCHW -> BHWC + elif debug_mode == "test_pattern": + # Use test pattern to clearly see warping effect + B, H, W, C = data.shape + test_pattern = self.create_test_pattern(H, W, device=data.device) + data = test_pattern # Replace input with test pattern + # Also show this as debug visualization + debug_vis = test_pattern + else: + # No specific debug visualization + debug_vis = torch.zeros_like(data) + else: + # For latent, just create a placeholder visualization + debug_vis = torch.zeros((1, 128, 128, 3), device=flow_field.device) + + # Process based on data type + if data_type == "IMAGE": + B, H, W, C = data.shape + device = data.device + + # Create coordinate grid in pixel space (0 to H-1, 0 to W-1) + y_coords = torch.arange(0, H, device=device).float() + x_coords = torch.arange(0, W, device=device).float() + grid_y, grid_x = torch.meshgrid(y_coords, x_coords, indexing="ij") + pixel_grid = torch.stack([grid_x, grid_y], dim=-1) # [H, W, 2] + + # Apply flow to pixel coordinates (forward warping) + # But we need backward warping for grid_sample, so we need to convert + # Flow is "where pixels move to" but we need "where to sample from" + + # Convert flow from BCHW [B, 2, H, W] to BHWC [B, H, W, 2] + flow_grid = scaled_flow.permute(0, 2, 3, 1).clone() # Clone to avoid in-place issues + + # Apply flow to destination coordinates + # dst_coords = pixel_grid + flow_grid + + # For grid_sample, we need normalized coords in [-1, 1] + # First, build a basic normalized grid + norm_y = 2.0 * torch.arange(H, device=device) / (H - 1) - 1.0 + norm_x = 2.0 * torch.arange(W, device=device) / (W - 1) - 1.0 + norm_grid_y, norm_grid_x = torch.meshgrid(norm_y, norm_x, indexing="ij") + identity_grid = torch.stack([norm_grid_x, norm_grid_y], dim=-1) # [H, W, 2] + + # Expand to batch dimension + identity_grid = identity_grid.expand(B, H, W, 2) + + # For grid_sample, we need to know where to sample from + # Flow tells us where pixels move to, so we reverse it + # Normalize flow to [-1, 1] range + flow_grid[..., 0] = flow_grid[..., 0] * (2.0 / (W - 1)) # x component + flow_grid[..., 1] = flow_grid[..., 1] * (2.0 / (H - 1)) # y component + + # Build the sampling grid by subtracting normalized flow from identity grid + sampling_grid = identity_grid - flow_grid + + if debug_mode == "grid" and data_type == "IMAGE": + # Visualize grid distortion + debug_img = torch.zeros((B, H, W, 3), device=device) + + # Normalize grid to [0, 1] for visualization + grid_vis = (sampling_grid + 1) / 2.0 + + # Assign x and y components to R and G channels + debug_img[..., 0] = grid_vis[..., 0] # R = x + debug_img[..., 1] = grid_vis[..., 1] # G = y + + # Draw grid lines + grid_size = 0.1 + grid_x = torch.abs(torch.frac(grid_vis[..., 0] / grid_size) - 0.5) < 0.02 + grid_y = torch.abs(torch.frac(grid_vis[..., 1] / grid_size) - 0.5) < 0.02 + debug_img[..., 2] = torch.logical_or(grid_x, grid_y).float() + + debug_vis = debug_img + + # Convert image from BHWC to BCHW for grid_sample + image_bchw = data.permute(0, 3, 1, 2) + + # Apply grid sampling + warped_image_bchw = F.grid_sample( + image_bchw, + sampling_grid, + mode=mode, + padding_mode=padding_mode, + align_corners=align_corners + ) + + # Convert back to BHWC + warped_image = warped_image_bchw.permute(0, 2, 3, 1) + + # For test pattern, return both the warped pattern and the original + if debug_mode == "test_pattern": + # Show side-by-side comparison + combined = torch.cat([data, warped_image], dim=2) # Concatenate along width + debug_vis = combined + + return (warped_image, debug_vis) + + elif data_type == "LATENT": + # Handle LATENT data type + if not isinstance(data, dict) or "samples" not in data: + raise ValueError("Expected latent input to be a dictionary with 'samples' key") + + samples = data["samples"] + B, C, H, W = samples.shape + device = samples.device + + # Resize flow to match latent dimensions if needed + if flow_field.shape[2:] != samples.shape[2:]: + scaled_flow = F.interpolate( + scaled_flow, + size=(H, W), + mode='bilinear', + align_corners=True + ) + + # Build the sampling grid + norm_y = 2.0 * torch.arange(H, device=device) / (H - 1) - 1.0 + norm_x = 2.0 * torch.arange(W, device=device) / (W - 1) - 1.0 + norm_grid_y, norm_grid_x = torch.meshgrid(norm_y, norm_x, indexing="ij") + identity_grid = torch.stack([norm_grid_x, norm_grid_y], dim=-1) # [H, W, 2] + + # Expand to batch dimension + identity_grid = identity_grid.expand(B, H, W, 2) + + # Convert flow from BCHW [B, 2, H, W] to BHWC [B, H, W, 2] + flow_grid = scaled_flow.permute(0, 2, 3, 1).clone() + + # Normalize flow to [-1, 1] range + flow_grid[..., 0] = flow_grid[..., 0] * (2.0 / (W - 1)) # x component + flow_grid[..., 1] = flow_grid[..., 1] * (2.0 / (H - 1)) # y component + + # Build the sampling grid by subtracting normalized flow from identity grid + sampling_grid = identity_grid - flow_grid + + # Apply grid sampling (latent is already in BCHW format) + warped_samples = F.grid_sample( + samples, + sampling_grid, + mode=mode, + padding_mode=padding_mode, + align_corners=align_corners + ) + + # Create new latent dict with warped samples + warped_latent = {k: v for k, v in data.items()} # Shallow copy + warped_latent["samples"] = warped_samples + + return (warped_latent, debug_vis) + + else: + raise ValueError(f"Unsupported data type: {data_type}") \ No newline at end of file diff --git a/node_wrappers/realtimenodes/controls/transform_controls.py b/node_wrappers/realtimenodes/controls/transform_controls.py new file mode 100644 index 0000000..06bc874 --- /dev/null +++ b/node_wrappers/realtimenodes/controls/transform_controls.py @@ -0,0 +1,327 @@ +""" +Transform control nodes for manipulating and creating transformation matrices. + +These nodes allow for precise control of affine transformations in real-time. +""" + +import torch +import math + +from ....src.realtimenodes.control_base import ControlNodeBase +from ....src.utils.transforms import ( + create_identity_matrix, + create_translation_matrix, + create_rotation_matrix, + create_scale_matrix, + compose_transforms +) + + +class TransformMatrixNode(ControlNodeBase): + """ + Node that maintains and manipulates an affine transformation matrix with precise control. + + Allows for precise mathematical operations on transformation matrices. + """ + + @classmethod + def INPUT_TYPES(cls): + inputs = super().INPUT_TYPES() + inputs["required"].update({ + "operation": ( + ["initialize", "translate", "rotate", "scale", "compose", "reset"], + {"default": "initialize"} + ), + "translate_x": ( + "FLOAT", + {"default": 0.0, "step": 0.01, "tooltip": "Translation in x direction (-1 to 1 range)"} + ), + "translate_y": ( + "FLOAT", + {"default": 0.0, "step": 0.01, "tooltip": "Translation in y direction (-1 to 1 range)"} + ), + "rotate_angle": ( + "FLOAT", + {"default": 0.0, "step": 0.1, "tooltip": "Rotation angle in degrees"} + ), + "scale_factor_x": ( + "FLOAT", + {"default": 1.0, "step": 0.01, "tooltip": "Scale factor in x direction"} + ), + "scale_factor_y": ( + "FLOAT", + {"default": 1.0, "step": 0.01, "tooltip": "Scale factor in y direction"} + ), + "center_x": ( + "FLOAT", + {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, + "tooltip": "X coordinate of transformation center (0-1 range)"} + ), + "center_y": ( + "FLOAT", + {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, + "tooltip": "Y coordinate of transformation center (0-1 range)"} + ), + }) + inputs["optional"] = { + "secondary_matrix": ( + "TRANSFORM_MATRIX", + {"tooltip": "Secondary matrix for composition operations"} + ) + } + return inputs + + RETURN_TYPES = ("TRANSFORM_MATRIX",) + RETURN_NAMES = ("transform_matrix",) + FUNCTION = "update" + CATEGORY = "Realtime Nodes/control/transform" + + def update( + self, + operation, + translate_x, + translate_y, + rotate_angle, + scale_factor_x, + scale_factor_y, + center_x, + center_y, + secondary_matrix=None, + always_execute=True, + unique_id=None + ): + """ + Update the transformation matrix based on the specified operation. + + Args: + operation: The operation to perform on the matrix + translate_x: X translation amount + translate_y: Y translation amount + rotate_angle: Rotation angle in degrees + scale_factor_x: X scale factor + scale_factor_y: Y scale factor + center_x: X coordinate of transformation center + center_y: Y coordinate of transformation center + secondary_matrix: Optional secondary matrix for composition + always_execute: Whether to always execute (from ControlNodeBase) + unique_id: Unique ID for state management (from ControlNodeBase) + + Returns: + Updated transformation matrix + """ + # Get state or create identity matrix + state = self.get_state({"matrix": None}, unique_id) + + # Default to CPU if we don't have a device yet + device = "cpu" + if state["matrix"] is not None: + device = state["matrix"].device + elif secondary_matrix is not None: + device = secondary_matrix.device + + current_matrix = state["matrix"] if state["matrix"] is not None else create_identity_matrix(device) + + # Perform operation + if operation == "initialize": + current_matrix = create_identity_matrix(device) + + elif operation == "translate": + t_matrix = create_translation_matrix(translate_x, translate_y, device) + current_matrix = compose_transforms(t_matrix, current_matrix) + + elif operation == "rotate": + r_matrix = create_rotation_matrix(rotate_angle, center_x, center_y, device) + current_matrix = compose_transforms(r_matrix, current_matrix) + + elif operation == "scale": + s_matrix = create_scale_matrix(scale_factor_x, scale_factor_y, center_x, center_y, device) + current_matrix = compose_transforms(s_matrix, current_matrix) + + elif operation == "compose" and secondary_matrix is not None: + current_matrix = compose_transforms(secondary_matrix, current_matrix) + + elif operation == "reset": + current_matrix = create_identity_matrix(device) + + # Update state + state["matrix"] = current_matrix + self.set_state(state, unique_id) + + return (current_matrix,) + + +class CameraControlNode(ControlNodeBase): + """ + High-level camera control with intuitive, independent parameters. + + Provides artist-friendly controls for camera manipulation in real-time. + """ + + @classmethod + def INPUT_TYPES(cls): + inputs = super().INPUT_TYPES() + inputs["required"].update({ + "pan_x": ( + "FLOAT", + {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01, + "tooltip": "Horizontal panning (-1.0 to 1.0)"} + ), + "pan_y": ( + "FLOAT", + {"default": 0.0, "min": -1.0, "max": 1.0, "step": 0.01, + "tooltip": "Vertical panning (-1.0 to 1.0)"} + ), + "zoom": ( + "FLOAT", + {"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.01, + "tooltip": "Zoom factor (1.0 = no zoom)"} + ), + "rotation": ( + "FLOAT", + {"default": 0.0, "min": -180.0, "max": 180.0, "step": 0.1, + "tooltip": "Rotation in degrees"} + ), + "center_x": ( + "FLOAT", + {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, + "tooltip": "X coordinate of rotation/zoom center"} + ), + "center_y": ( + "FLOAT", + {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, + "tooltip": "Y coordinate of rotation/zoom center"} + ), + "reset": ( + "BOOLEAN", + {"default": False, "tooltip": "Reset camera to default state"} + ), + "accumulate": ( + "BOOLEAN", + {"default": False, "tooltip": "Accumulate transformations over time (vs. absolute values)"} + ), + }) + inputs["optional"] = { + "base_matrix": ( + "TRANSFORM_MATRIX", + {"tooltip": "Starting transform matrix (optional)"} + ) + } + return inputs + + RETURN_TYPES = ("TRANSFORM_MATRIX",) + RETURN_NAMES = ("transform_matrix",) + FUNCTION = "update" + CATEGORY = "Realtime Nodes/control/transform" + + def update( + self, + pan_x, + pan_y, + zoom, + rotation, + center_x, + center_y, + reset, + accumulate, + base_matrix=None, + always_execute=True, + unique_id=None + ): + """ + Update the camera transformation matrix. + + Args: + pan_x: Horizontal panning amount + pan_y: Vertical panning amount + zoom: Zoom factor + rotation: Rotation angle in degrees + center_x: X coordinate of rotation/zoom center + center_y: Y coordinate of rotation/zoom center + reset: Whether to reset the camera + accumulate: Whether to accumulate transforms (vs. absolute values) + base_matrix: Optional base matrix to start from + always_execute: Whether to always execute (from ControlNodeBase) + unique_id: Unique ID for state management (from ControlNodeBase) + + Returns: + Updated camera transformation matrix + """ + # Get state or initialize + state = self.get_state({ + "matrix": None, # Current transform matrix + "last_params": { # Last used parameters + "pan_x": 0.0, + "pan_y": 0.0, + "zoom": 1.0, + "rotation": 0.0 + } + }, unique_id) + + device = "cpu" + + # Handle reset + if reset: + state["matrix"] = None + state["last_params"] = { + "pan_x": 0.0, + "pan_y": 0.0, + "zoom": 1.0, + "rotation": 0.0 + } + + # Start with base_matrix or identity + if base_matrix is not None: + current_matrix = base_matrix + device = base_matrix.device + elif not accumulate: + # For non-accumulate mode, always start fresh + current_matrix = create_identity_matrix(device) + elif state["matrix"] is not None: + current_matrix = state["matrix"] + device = current_matrix.device + else: + current_matrix = create_identity_matrix(device) + + # For accumulate mode, calculate the delta from last parameters + if accumulate and state["last_params"] is not None: + # Calculate deltas + delta_pan_x = pan_x - state["last_params"]["pan_x"] + delta_pan_y = pan_y - state["last_params"]["pan_y"] + delta_zoom = zoom / state["last_params"]["zoom"] if state["last_params"]["zoom"] != 0 else zoom + delta_rotation = rotation - state["last_params"]["rotation"] + + # Use deltas instead of absolute values + pan_x = delta_pan_x + pan_y = delta_pan_y + zoom = delta_zoom + rotation = delta_rotation + + # Apply transforms in sequence: first scale (zoom), then rotate, then translate (pan) + # This order ensures consistent behavior + + # 1. Apply zoom + if zoom != 1.0: + zoom_matrix = create_scale_matrix(zoom, zoom, center_x, center_y, device) + current_matrix = compose_transforms(zoom_matrix, current_matrix) + + # 2. Apply rotation + if rotation != 0.0: + rotation_matrix = create_rotation_matrix(rotation, center_x, center_y, device) + current_matrix = compose_transforms(rotation_matrix, current_matrix) + + # 3. Apply panning + if pan_x != 0.0 or pan_y != 0.0: + translation_matrix = create_translation_matrix(pan_x, pan_y, device) + current_matrix = compose_transforms(translation_matrix, current_matrix) + + # Update state with current parameters and matrix + state["matrix"] = current_matrix + state["last_params"] = { + "pan_x": pan_x if not accumulate else pan_x + state["last_params"]["pan_x"], + "pan_y": pan_y if not accumulate else pan_y + state["last_params"]["pan_y"], + "zoom": zoom if not accumulate else zoom * state["last_params"]["zoom"], + "rotation": rotation if not accumulate else rotation + state["last_params"]["rotation"] + } + self.set_state(state, unique_id) + + return (current_matrix,) \ No newline at end of file diff --git a/node_wrappers/realtimenodes/controls/transform_warpers.py b/node_wrappers/realtimenodes/controls/transform_warpers.py new file mode 100644 index 0000000..994b350 --- /dev/null +++ b/node_wrappers/realtimenodes/controls/transform_warpers.py @@ -0,0 +1,152 @@ +""" +Transform warping nodes for applying transformation matrices to images and latents. + +These nodes apply affine transformations to various media types in real-time. +""" + +import torch +import torch.nn.functional as F + +from ....src.realtimenodes.control_base import ControlNodeBase +from ....src.utils.warp import warp_image, warp_latent + +#TODO add depth aware throughout + +class MatrixTransformImageNode(ControlNodeBase): + """ + Apply affine transformation matrix to images. + + Warps images using a provided transformation matrix. + """ + + @classmethod + def INPUT_TYPES(cls): + inputs = super().INPUT_TYPES() + inputs["required"].update({ + "image": ("IMAGE",), + "transform_matrix": ("TRANSFORM_MATRIX",), + "mode": ( + ["bilinear", "nearest", "bicubic"], + {"default": "bilinear", "tooltip": "Interpolation mode for sampling"} + ), + "padding_mode": ( + ["zeros", "border", "reflection"], + {"default": "border", "tooltip": "How to handle pixels outside the image boundaries"} + ), + "align_corners": ( + "BOOLEAN", + {"default": True, "tooltip": "Whether to align corners when sampling"} + ), + }) + return inputs + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("transformed_image",) + FUNCTION = "update" + CATEGORY = "Realtime Nodes/control/transform" + + def update( + self, + image, + transform_matrix, + mode="bilinear", + padding_mode="border", + align_corners=True, + always_execute=True, + unique_id=None + ): + """ + Apply transformation matrix to an image. + + Args: + image: Input image tensor [B, H, W, C] + transform_matrix: 3x3 transformation matrix + mode: Interpolation mode + padding_mode: Padding mode + align_corners: Whether to align corners in sampling + always_execute: Whether to always execute (from ControlNodeBase) + unique_id: Unique ID for state management (from ControlNodeBase) + + Returns: + Warped image + """ + # Apply warping to image using utility function + warped_image = warp_image( + image, + transform_matrix, + mode=mode, + padding_mode=padding_mode, + align_corners=align_corners + ) + + return (warped_image,) + + +class MatrixTransformLatentNode(ControlNodeBase): + """ + Apply affine transformation matrix to latent space. + + Warps latent tensors using a provided transformation matrix. + """ + + @classmethod + def INPUT_TYPES(cls): + inputs = super().INPUT_TYPES() + inputs["required"].update({ + "latent": ("LATENT",), + "transform_matrix": ("TRANSFORM_MATRIX",), + "mode": ( + ["bilinear", "nearest", "bicubic"], + {"default": "bilinear", "tooltip": "Interpolation mode for sampling"} + ), + "padding_mode": ( + ["zeros", "border", "reflection"], + {"default": "border", "tooltip": "How to handle values outside the latent boundaries"} + ), + "align_corners": ( + "BOOLEAN", + {"default": True, "tooltip": "Whether to align corners when sampling"} + ), + }) + return inputs + + RETURN_TYPES = ("LATENT",) + RETURN_NAMES = ("transformed_latent",) + FUNCTION = "update" + CATEGORY = "Realtime Nodes/control/transform" + + def update( + self, + latent, + transform_matrix, + mode="bilinear", + padding_mode="border", + align_corners=True, + always_execute=True, + unique_id=None + ): + """ + Apply transformation matrix to latent space. + + Args: + latent: Input latent dict with 'samples' tensor [B, C, H, W] + transform_matrix: 3x3 transformation matrix + mode: Interpolation mode + padding_mode: Padding mode + align_corners: Whether to align corners in sampling + always_execute: Whether to always execute (from ControlNodeBase) + unique_id: Unique ID for state management (from ControlNodeBase) + + Returns: + Warped latent + """ + # Apply warping to latent using utility function + warped_latent = warp_latent( + latent, + transform_matrix, + mode=mode, + padding_mode=padding_mode, + align_corners=align_corners + ) + + return (warped_latent,) \ No newline at end of file diff --git a/src/utils/__init__.py b/src/utils/__init__.py index 22c9785..4bf409a 100644 --- a/src/utils/__init__.py +++ b/src/utils/__init__.py @@ -6,7 +6,7 @@ Unified utilities package for ComfyUI RealTimeNodes. from .general import AlwaysEqualProxy # Image utilities -from .image import convert_to_cv2, convert_to_tensor, create_mask_from_points +from .image import convert_to_cv2, convert_to_tensor, create_mask_from_points, flow_to_rgb, gaussian_blur_2d # Math utilities from .math import ( @@ -18,6 +18,31 @@ from .math import ( # Timing utilities from .timing import TimestampProvider +# Similar image filter +from .similar_image_filter import SimilarImageFilter + +# Transform utilities +from .transforms import ( + create_identity_matrix, + create_translation_matrix, + create_rotation_matrix, + create_scale_matrix, + compose_transforms, + normalize_homogeneous_coordinates, + apply_transform +) + +# Warp utilities +from .warp import ( + create_sampling_grid, + transform_grid, + warp_image, + warp_latent +) + +# Real-time flownets +from .realtime_flownets import RealTimeFlowNet + # For backward compatibility, expose the original names # This allows existing code to continue working while we transition image_to_cv2 = convert_to_cv2 diff --git a/src/utils/transforms.py b/src/utils/transforms.py new file mode 100644 index 0000000..957eac5 --- /dev/null +++ b/src/utils/transforms.py @@ -0,0 +1,174 @@ +""" +Transform utilities for manipulating affine transformation matrices. + +Contains functions for creating and composing transformation matrices. +""" + +import torch +import math + + +def create_identity_matrix(device="cpu"): + """ + Create a 3x3 identity transformation matrix. + + Args: + device: The torch device to place the matrix on + + Returns: + 3x3 identity matrix as torch.Tensor + """ + return torch.eye(3, device=device) + + +def create_translation_matrix(dx, dy, device="cpu"): + """ + Create a 3x3 translation matrix. + + Args: + dx: Translation in x direction + dy: Translation in y direction + device: The torch device to place the matrix on + + Returns: + 3x3 translation matrix as torch.Tensor + """ + matrix = torch.eye(3, device=device) + matrix[0, 2] = dx + matrix[1, 2] = dy + return matrix + + +def create_rotation_matrix(angle_degrees, center_x=0.5, center_y=0.5, device="cpu"): + """ + Create a 3x3 rotation matrix that rotates around the specified center point. + + Args: + angle_degrees: Rotation angle in degrees + center_x: X coordinate of rotation center (0-1 range) + center_y: Y coordinate of rotation center (0-1 range) + device: The torch device to place the matrix on + + Returns: + 3x3 rotation matrix as torch.Tensor + """ + angle_rad = math.radians(angle_degrees) + cos_theta = math.cos(angle_rad) + sin_theta = math.sin(angle_rad) + + # Convert center from 0-1 range to -1 to 1 range + center_x = (center_x * 2) - 1 + center_y = (center_y * 2) - 1 + + # 1. Translate center to origin + t1 = create_translation_matrix(-center_x, -center_y, device) + + # 2. Rotate around origin + r = torch.eye(3, device=device) + r[0, 0] = cos_theta + r[0, 1] = -sin_theta + r[1, 0] = sin_theta + r[1, 1] = cos_theta + + # 3. Translate back + t2 = create_translation_matrix(center_x, center_y, device) + + # Compose transformations: t2 @ r @ t1 + return torch.matmul(t2, torch.matmul(r, t1)) + + +def create_scale_matrix(scale_x, scale_y, center_x=0.5, center_y=0.5, device="cpu"): + """ + Create a 3x3 scaling matrix that scales around the specified center point. + + Args: + scale_x: Scale factor in x direction + scale_y: Scale factor in y direction + center_x: X coordinate of scaling center (0-1 range) + center_y: Y coordinate of scaling center (0-1 range) + device: The torch device to place the matrix on + + Returns: + 3x3 scaling matrix as torch.Tensor + """ + # Convert center from 0-1 range to -1 to 1 range + center_x = (center_x * 2) - 1 + center_y = (center_y * 2) - 1 + + # 1. Translate center to origin + t1 = create_translation_matrix(-center_x, -center_y, device) + + # 2. Scale around origin + s = torch.eye(3, device=device) + s[0, 0] = scale_x + s[1, 1] = scale_y + + # 3. Translate back + t2 = create_translation_matrix(center_x, center_y, device) + + # Compose transformations: t2 @ s @ t1 + return torch.matmul(t2, torch.matmul(s, t1)) + + +def compose_transforms(transform_a, transform_b): + """ + Compose two transformation matrices. The result is equivalent to applying + transform_b first, then transform_a. + + Args: + transform_a: First transformation matrix + transform_b: Second transformation matrix + + Returns: + Composed transformation matrix + """ + return torch.matmul(transform_a, transform_b) + + +def normalize_homogeneous_coordinates(coords): + """ + Normalize homogeneous coordinates by dividing by the last component. + + Args: + coords: Tensor of homogeneous coordinates [..., 3] + + Returns: + Normalized coordinates + """ + # Extract last component + w = coords[..., 2:3] + + # Handle zeros (avoid division by zero) + mask = w != 0 + + # Create a copy of the coordinates + normalized = coords.clone() + + # Apply normalization only to non-zero denominator elements + # We need to expand the mask to match the coords dimension + for i in range(coords.shape[-1]): + normalized[..., i] = torch.where(mask.squeeze(-1), coords[..., i] / w.squeeze(-1), coords[..., i]) + + return normalized + + +def apply_transform(points, transform_matrix): + """ + Apply a transformation matrix to a set of points. + + Args: + points: Tensor of points [..., 2] + transform_matrix: 3x3 transformation matrix + + Returns: + Transformed points + """ + # Convert to homogeneous coordinates + homogeneous = torch.cat([points, torch.ones_like(points[..., :1])], dim=-1) + + # Apply transformation + transformed = torch.matmul(homogeneous, transform_matrix.T) + + # Normalize and return x, y coordinates + normalized = normalize_homogeneous_coordinates(transformed) + return normalized[..., :2] \ No newline at end of file diff --git a/src/utils/warp.py b/src/utils/warp.py new file mode 100644 index 0000000..59f3d86 --- /dev/null +++ b/src/utils/warp.py @@ -0,0 +1,181 @@ +""" +Warping utilities for transforming images and latents using transformation matrices. + +Contains functions for grid-based warping operations. +""" + +import torch +import torch.nn.functional as F +from .transforms import normalize_homogeneous_coordinates + + +def create_sampling_grid(height, width, device="cpu"): + """ + Create a normalized sampling grid for an image of the given dimensions. + + Args: + height: Image height + width: Image width + device: The torch device to place the grid on + + Returns: + Normalized grid of shape [height, width, 2] with values in range [-1, 1] + """ + # Create coordinate grid in pixel space + y_coords = torch.linspace(0, height - 1, height, device=device) + x_coords = torch.linspace(0, width - 1, width, device=device) + grid_y, grid_x = torch.meshgrid(y_coords, x_coords, indexing="ij") + + # Convert to normalized coordinates in range [-1, 1] + grid_x = 2.0 * grid_x / (width - 1) - 1.0 + grid_y = 2.0 * grid_y / (height - 1) - 1.0 + + # Stack into grid of shape [height, width, 2] + grid = torch.stack([grid_x, grid_y], dim=-1) + + return grid + + +def transform_grid(grid, transform_matrix): + """ + Apply a transformation matrix to a sampling grid. + + The transformation matrix is applied in reverse (inverse) since we're creating + a sampling grid - we need to know where to sample from in the source image. + + Args: + grid: Sampling grid of shape [..., 2] + transform_matrix: 3x3 transformation matrix + + Returns: + Transformed grid of same shape as input + """ + # Get grid shape for reshaping later + original_shape = grid.shape + device = grid.device + + # Flatten the grid to make computations easier + # Reshape to [N, 2] where N is the product of all dimensions except the last + flat_shape = (-1, 2) + grid_flat = grid.reshape(flat_shape) + + # Convert to homogeneous coordinates + ones = torch.ones(grid_flat.shape[0], 1, device=device) + homogeneous_grid = torch.cat([grid_flat, ones], dim=-1) + + # Compute inverse for sampling (where to get pixels from) + try: + inverse_matrix = torch.inverse(transform_matrix) + except RuntimeError: + # Handle singular matrix by adding small epsilon + epsilon = 1e-6 + identity = torch.eye(3, device=device) * epsilon + transform_matrix_reg = transform_matrix + identity + inverse_matrix = torch.inverse(transform_matrix_reg) + + # Apply transformation + transformed = torch.matmul(homogeneous_grid, inverse_matrix.T) + + # Normalize homogeneous coordinates (divide by w component) + w = transformed[:, 2:3] + mask = w != 0 + normalized = torch.zeros_like(transformed) + + # Apply normalization to x, y, z components where w != 0 + normalized[:, 0] = torch.where(mask.squeeze(-1), transformed[:, 0] / w.squeeze(-1), transformed[:, 0]) + normalized[:, 1] = torch.where(mask.squeeze(-1), transformed[:, 1] / w.squeeze(-1), transformed[:, 1]) + normalized[:, 2] = torch.where(mask.squeeze(-1), transformed[:, 2] / w.squeeze(-1), transformed[:, 2]) + + # Get x, y components and reshape back to original grid shape + transformed_grid = normalized[:, :2].reshape(original_shape) + + return transformed_grid + + +def warp_image(image, transform_matrix, mode="bilinear", padding_mode="zeros", align_corners=True): + """ + Warp an image using a transformation matrix. + + Args: + image: Input image tensor of shape [B, H, W, C] + transform_matrix: 3x3 transformation matrix + mode: Interpolation mode ('bilinear', 'nearest', or 'bicubic') + padding_mode: Padding mode ('zeros', 'border', or 'reflection') + align_corners: Whether to align corners in the grid_sample operation + + Returns: + Warped image of same shape as input + """ + # Get image dimensions + B, H, W, C = image.shape + device = image.device + + # Create sampling grid + grid = create_sampling_grid(H, W, device) + + # Apply transformation to grid + warped_grid = transform_grid(grid, transform_matrix) + + # Expand grid to batch dimension + warped_grid = warped_grid.expand(B, -1, -1, -1) + + # Reorder image for grid_sample: [B, C, H, W] + image_bchw = image.permute(0, 3, 1, 2) + + # Apply grid sampling + warped_image = F.grid_sample( + image_bchw, + warped_grid, + mode=mode, + padding_mode=padding_mode, + align_corners=align_corners + ) + + # Reorder back to [B, H, W, C] + warped_image = warped_image.permute(0, 2, 3, 1) + + return warped_image + + +def warp_latent(latent, transform_matrix, mode="bilinear", padding_mode="zeros", align_corners=True): + """ + Warp a latent tensor using a transformation matrix. + + Args: + latent: Input latent dict with 'samples' tensor of shape [B, C, H, W] + transform_matrix: 3x3 transformation matrix + mode: Interpolation mode ('bilinear', 'nearest', or 'bicubic') + padding_mode: Padding mode ('zeros', 'border', or 'reflection') + align_corners: Whether to align corners in the grid_sample operation + + Returns: + Warped latent in same format as input + """ + # Get latent samples + samples = latent["samples"] + B, C, H, W = samples.shape + device = samples.device + + # Create sampling grid + grid = create_sampling_grid(H, W, device) + + # Apply transformation to grid + warped_grid = transform_grid(grid, transform_matrix) + + # Expand grid to batch dimension + warped_grid = warped_grid.expand(B, -1, -1, -1) + + # Apply grid sampling (latent is already in [B, C, H, W] format) + warped_samples = F.grid_sample( + samples, + warped_grid, + mode=mode, + padding_mode=padding_mode, + align_corners=align_corners + ) + + # Create new latent dict with warped samples + warped_latent = latent.copy() + warped_latent["samples"] = warped_samples + + return warped_latent \ No newline at end of file