From faf67c4ddc8f34499a993b37c4e72ea9a97214c6 Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Tue, 4 Mar 2025 15:27:20 -0500 Subject: [PATCH] generator prototype --- README.MD | 3 + __init__.py | 6 +- misc_nodes.py | 15 +++ particle_nodes.py | 244 ++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 267 insertions(+), 1 deletion(-) create mode 100644 particle_nodes.py diff --git a/README.MD b/README.MD index 0142302..4a8d584 100644 --- a/README.MD +++ b/README.MD @@ -21,6 +21,9 @@ The intention for this repository is to build a suite of nodes that can be used - **IntegerMotionController**: Integer-based motion control for discrete value animations. - **ROINode**: Region of Interest node for motion tracking and control. +### Particle Effects 🌟 +- **TemporalParticleDepthNode**: Generates a temporally consistent particle depth map with a controllable origin point. Can be linked to hand tracking for interactive particle effects. + ### Utility Controls 🛠️ - **FPSMonitor**: Generates an FPS overlay as an image and mask, useful for monitoring performance. - **QuickShapeMask**: Rapidly generate shape masks (circle, square) with customizable dimensions. diff --git a/__init__.py b/__init__.py index 4629c9c..a4786be 100644 --- a/__init__.py +++ b/__init__.py @@ -10,10 +10,12 @@ from .misc_nodes import( TextRenderer, QuickShapeMask, MultilineText, - LoadImageFromPath_ + LoadImageFromPath_, + Noop ) from .media_pipe_nodes import HandTrackingNode, HandMaskNode from .controls.mask_controls import RepulsiveMaskNode, ResizeMaskNode +from .particle_nodes import TemporalParticleDepthNode import re @@ -38,12 +40,14 @@ NODE_CLASS_MAPPINGS = { "FastWebcamCapture": FastWebcamCapture, "MultilineText": MultilineText, "LoadImageFromPath_": LoadImageFromPath_, + "Noop": Noop, "HandTrackingNode": HandTrackingNode, "HandMaskNode": HandMaskNode, #"RepulsiveMaskNode": RepulsiveMaskNode, "ResizeMaskNode": ResizeMaskNode, "StateResetNode": StateResetNode, "StateTestNode": StateTestNode, + "TemporalParticleDepthNode": TemporalParticleDepthNode, } diff --git a/misc_nodes.py b/misc_nodes.py index c7aa5c6..436ca7b 100644 --- a/misc_nodes.py +++ b/misc_nodes.py @@ -573,3 +573,18 @@ class LoadImageFromPath_: return (image, mask) +class Noop: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "any": (AlwaysEqualProxy("*"),), + } + } + + RETURN_TYPES = () + FUNCTION = "noop" + CATEGORY = "utils" + + def noop(self, any): + pass diff --git a/particle_nodes.py b/particle_nodes.py new file mode 100644 index 0000000..5b1512b --- /dev/null +++ b/particle_nodes.py @@ -0,0 +1,244 @@ +import torch +import numpy as np +from .base.control_base import ControlNodeBase + +class TemporalParticleDepthNode(ControlNodeBase): + """ + A node that generates temporally consistent particle depth maps with a controllable origin point. + """ + + @classmethod + def INPUT_TYPES(cls): + inputs = super().INPUT_TYPES() + inputs["required"].update({ + "origin_x": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "display": "slider", + "tooltip": "X coordinate of the origin point (0-1)" + }), + "origin_y": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "display": "slider", + "tooltip": "Y coordinate of the origin point (0-1)" + }), + "width": ("INT", { + "default": 512, + "min": 64, + "max": 2048, + "step": 8, + "tooltip": "Width of the output depth map" + }), + "height": ("INT", { + "default": 512, + "min": 64, + "max": 2048, + "step": 8, + "tooltip": "Height of the output depth map" + }), + "num_particles": ("INT", { + "default": 100, + "min": 10, + "max": 1000, + "step": 10, + "tooltip": "Number of particles to simulate" + }), + "particle_size": ("FLOAT", { + "default": 0.03, + "min": 0.001, + "max": 0.2, + "step": 0.001, + "tooltip": "Size of each particle (relative to image size)" + }), + }) + inputs["optional"] = { + "hand_data": ("HAND_DATA", { + "tooltip": "Optional hand tracking data to use for origin point" + }), + "hand_keypoint": (["palm", "thumb_tip", "index_tip", "middle_tip", "ring_tip", "pinky_tip"], { + "default": "index_tip", + "tooltip": "Which hand keypoint to use as origin (if hand_data is provided)" + }), + } + return inputs + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "update" + CATEGORY = "image/generators" + + def __init__(self): + super().__init__() + self.generator = None + self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + + def update(self, origin_x, origin_y, width, height, num_particles, particle_size, always_execute=True, hand_data=None, hand_keypoint="index_tip"): + # Initialize generator if needed or if parameters changed + if self.generator is None: + self.generator = TemporalParticleDepthGenerator( + batch_size=1, + height=height, + width=width, + num_particles=num_particles, + initial_origin_x=origin_x, + initial_origin_y=origin_y, + particle_size=particle_size, + device=self.device + ) + elif (self.generator.shape[1] != height or + self.generator.shape[2] != width or + self.generator.num_particles != num_particles): + # Reinitialize if dimensions or particle count changed + self.generator = TemporalParticleDepthGenerator( + batch_size=1, + height=height, + width=width, + num_particles=num_particles, + initial_origin_x=origin_x, + initial_origin_y=origin_y, + particle_size=particle_size, + device=self.device + ) + + # Use hand data if provided + if hand_data is not None and len(hand_data) > 0: + # Extract the first hand's data + hand_info = hand_data[0] + + # Check if hands are present + if hand_info["hands_present"]: + # Determine which hand to use (prefer right hand if available) + hand_landmarks = None + if hand_info["right_hand"] is not None: + hand_landmarks = hand_info["right_hand"] + elif hand_info["left_hand"] is not None: + hand_landmarks = hand_info["left_hand"] + + if hand_landmarks is not None: + # Map keypoint to index + keypoint_map = { + "palm": 0, # Center of palm + "thumb_tip": 4, # Thumb tip + "index_tip": 8, # Index finger tip + "middle_tip": 12, # Middle finger tip + "ring_tip": 16, # Ring finger tip + "pinky_tip": 20 # Pinky tip + } + + if hand_keypoint in keypoint_map: + keypoint_idx = keypoint_map[hand_keypoint] + # Get normalized coordinates (already in 0-1 range) + if keypoint_idx < len(hand_landmarks): + # X and Y coordinates are at indices 0 and 1 + origin_x = hand_landmarks[keypoint_idx][0] + origin_y = hand_landmarks[keypoint_idx][1] + + # Generate depth map with current origin + depth_map = self.generator.update( + new_origin_x=origin_x, + new_origin_y=origin_y, + particle_size=particle_size + ) + + # ComfyUI already uses BHWC format, so no permute needed + return (depth_map,) + + +class TemporalParticleDepthGenerator: + def __init__(self, + batch_size: int, + height: int, + width: int, + num_particles: int = 100, + initial_origin_x: float = 0.5, + initial_origin_y: float = 0.5, + particle_size: float = 0.03, + device: str = 'cuda' if torch.cuda.is_available() else 'cpu'): + """ + Initialize a temporally consistent particle depth map generator. + """ + self.device = torch.device(device) + self.shape = (batch_size, height, width) + self.particle_size = particle_size + + # Coordinate grid (static) + y = torch.linspace(0, 1, height, device=device) + x = torch.linspace(0, 1, width, device=device) + self.y_grid, self.x_grid = torch.meshgrid(y, x, indexing='ij') + + # Persistent particle properties + self.num_particles = num_particles + self.angles = torch.rand(num_particles, device=device) * 2 * torch.pi + self.speeds = torch.rand(num_particles, device=device) * 0.02 + 0.005 # Vary speed slightly + self.distances = torch.zeros(num_particles, device=device) # Start at origin + + # Origin can be updated later + self.origin_x = torch.tensor(initial_origin_x, device=device) + self.origin_y = torch.tensor(initial_origin_y, device=device) + + # For respawning particles that go off-screen + self.active = torch.ones(num_particles, dtype=torch.bool, device=device) + + def update(self, + new_origin_x: float = None, + new_origin_y: float = None, + particle_size: float = None) -> torch.Tensor: + """ + Update particle positions and generate next depth map frame. + + Args: + new_origin_x (float, optional): Update origin X (0-1) + new_origin_y (float, optional): Update origin Y (0-1) + particle_size (float, optional): Update particle size + + Returns: + torch.Tensor: Depth map (B, H, W, 3) + """ + # Update origin if provided + if new_origin_x is not None: + self.origin_x = torch.tensor(new_origin_x, device=self.device) + if new_origin_y is not None: + self.origin_y = torch.tensor(new_origin_y, device=self.device) + if particle_size is not None: + self.particle_size = particle_size + + # Update distances + self.distances = self.distances + self.speeds + + # Check which particles need respawning + off_screen = self.distances > 1.414 # Beyond sqrt(2), the diagonal + if off_screen.any(): + num_respawn = off_screen.sum() + self.distances[off_screen] = 0.0 + # Optional: Randomize angles for respawned particles + self.angles[off_screen] = torch.rand(num_respawn, device=self.device) * 2 * torch.pi + self.speeds[off_screen] = torch.rand(num_respawn, device=self.device) * 0.02 + 0.005 + + # Calculate current particle positions + particle_x = self.origin_x + torch.cos(self.angles) * self.distances + particle_y = self.origin_y + torch.sin(self.angles) * self.distances + + # Generate depth map + depth = torch.zeros(self.shape, device=self.device) + for i in range(self.num_particles): + # Skip if particle is inactive (future feature) + if not self.active[i]: + continue + + # Distance from particle center + dist = torch.sqrt((self.x_grid - particle_x[i])**2 + + (self.y_grid - particle_y[i])**2) + + # Gaussian particle shape + particle = torch.exp(-(dist**2) / (2 * self.particle_size**2)) + + # Depth based on distance from origin + depth[0] = torch.maximum(depth[0], particle * self.distances[i]) + + # Normalize and expand to BHWC + depth = depth / (depth.max() + 1e-6) + return depth.unsqueeze(-1).expand(-1, -1, -1, 3).clamp(0, 1) \ No newline at end of file