feat: setstate getstate
This primarily adds two new nodes -GetState -SetState It allows for state to be persisted and passed from one workflow run to another between get and set nodes. Additionally, there is an feedback effect node for experimentation.
This commit is contained in:
RyanOnTheInside
+37
-16
@@ -5,7 +5,7 @@ import re
|
||||
import traceback
|
||||
|
||||
# Directory containing node wrapper modules relative to this __init__.py
|
||||
WRAPPER_DIR = "node_wrappers.realtimenodes"
|
||||
WRAPPER_DIR = "node_wrappers"
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
|
||||
@@ -13,7 +13,7 @@ print("\033[93m" + "[ComfyUI_RealtimeNodes] Loading nodes...")
|
||||
|
||||
def load_nodes_from_directory(package_path):
|
||||
"""Dynamically loads nodes from Python files in a directory."""
|
||||
loaded_nodes = {}
|
||||
loaded_nodes = {}
|
||||
loaded_display_names = {}
|
||||
|
||||
# Calculate absolute path to the package directory
|
||||
@@ -61,22 +61,43 @@ def load_nodes_from_directory(package_path):
|
||||
|
||||
# --- Main Loading Logic ---
|
||||
|
||||
# Define the subdirectories within WRAPPER_DIR to scan
|
||||
# Order might matter if there are dependencies between modules, though ideally there shouldn't be
|
||||
subdirs_to_scan = [
|
||||
"controls",
|
||||
"media",
|
||||
"utils",
|
||||
"specialized" # Add other subdirs like 'specialized' if needed
|
||||
]
|
||||
# Get base path for the node_wrappers directory
|
||||
base_dir_path = os.path.join(os.path.dirname(__file__), WRAPPER_DIR)
|
||||
|
||||
# Iterate through specified subdirectories and load nodes
|
||||
for subdir in subdirs_to_scan:
|
||||
dir_path = f"{WRAPPER_DIR}.{subdir}"
|
||||
nodes, display_names = load_nodes_from_directory(dir_path)
|
||||
NODE_CLASS_MAPPINGS.update(nodes)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(display_names)
|
||||
# Recursively scan all directories and subdirectories
|
||||
def scan_for_modules(base_path, base_package):
|
||||
"""Recursively scan directories for module files"""
|
||||
try:
|
||||
for item in os.listdir(base_path):
|
||||
item_path = os.path.join(base_path, item)
|
||||
|
||||
# Skip __pycache__ and files/dirs starting with __
|
||||
if item.startswith('__'):
|
||||
continue
|
||||
|
||||
if os.path.isdir(item_path):
|
||||
# Handle subdirectory
|
||||
subdir_package = f"{base_package}.{item}"
|
||||
|
||||
# First load .py files in this directory
|
||||
try:
|
||||
nodes, display_names = load_nodes_from_directory(subdir_package)
|
||||
NODE_CLASS_MAPPINGS.update(nodes)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(display_names)
|
||||
except Exception as e:
|
||||
print(f"\033[91m" + f"[ComfyUI_RealtimeNodes] Error processing directory {subdir_package}: {e}")
|
||||
traceback.print_exc()
|
||||
|
||||
# Then recursively process subdirectories
|
||||
subdir_path = os.path.join(base_path, item)
|
||||
scan_for_modules(subdir_path, subdir_package)
|
||||
except Exception as e:
|
||||
print(f"\033[91m" + f"[ComfyUI_RealtimeNodes] Error scanning directory {base_path}: {e}")
|
||||
traceback.print_exc()
|
||||
|
||||
# Start recursive scanning from the base directory
|
||||
print(f"\033[92m" + f"[ComfyUI_RealtimeNodes] Scanning node_wrappers directory for modules...")
|
||||
scan_for_modules(base_dir_path, WRAPPER_DIR)
|
||||
|
||||
suffix = " 🕒🅡🅣🅝"
|
||||
|
||||
|
||||
@@ -1,4 +1,6 @@
|
||||
from ....src.realtimenodes.control_base import ControlNodeBase
|
||||
import copy
|
||||
from ....src.utils.general_utils import AlwaysEqualProxy
|
||||
|
||||
class StateResetNode(ControlNodeBase):
|
||||
"""Node that resets all control node states when triggered"""
|
||||
@@ -21,12 +23,8 @@ class StateResetNode(ControlNodeBase):
|
||||
CATEGORY = "real-time/control/utility"
|
||||
|
||||
def update(self, trigger, always_execute=True):
|
||||
print(f"\n=== StateResetNode UPDATE - node_id: {self.node_id} ===")
|
||||
print(f"States before potential reset: {self.state_manager._states}")
|
||||
if trigger:
|
||||
print("RESETTING ALL STATES")
|
||||
self.state_manager.clear_all_states()
|
||||
print(f"States after reset: {self.state_manager._states}")
|
||||
return (True,)
|
||||
return (False,)
|
||||
|
||||
@@ -52,19 +50,87 @@ class StateTestNode(ControlNodeBase):
|
||||
CATEGORY = "real-time/control/utility"
|
||||
|
||||
def update(self, increment, always_execute=True):
|
||||
print(f"\n=== StateTestNode UPDATE - node_id: {self.node_id} ===")
|
||||
print(f"All states before get: {self.state_manager._states}")
|
||||
|
||||
state = self.get_state({
|
||||
"counter": 0
|
||||
})
|
||||
print(f"Retrieved state: {state}")
|
||||
|
||||
state["counter"] += increment
|
||||
print(f"Updated state before save: {state}")
|
||||
|
||||
|
||||
self.set_state(state)
|
||||
print(f"All states after save: {self.state_manager._states}")
|
||||
|
||||
return (state["counter"],)
|
||||
|
||||
|
||||
|
||||
class GetStateNode(ControlNodeBase):
|
||||
"""
|
||||
Node that retrieves a value from the global state using a user-specified key.
|
||||
"""
|
||||
CATEGORY = "utils"
|
||||
RETURN_TYPES = (AlwaysEqualProxy("*"),)
|
||||
RETURN_NAMES = ("value",)
|
||||
FUNCTION = "update"
|
||||
DESCRIPTION = "Retrieve a value from the global state using the given key. If the key is not found, return the default value."
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
inputs = super().INPUT_TYPES()
|
||||
inputs["required"].update({
|
||||
"key": ("STRING", {"default": "default_key", "tooltip": "The key to retrieve the value from. If not provided, the default value will be returned."}),
|
||||
"default_value": (AlwaysEqualProxy("*"), {"tooltip": "The value to return if the key is not found."}),
|
||||
"use_default": ("BOOLEAN", {"default": False, "tooltip": "If True, the default value will be returned if the key is not found."}),
|
||||
})
|
||||
return inputs
|
||||
|
||||
def update(self, key: str, default_value, use_default: bool, always_execute=True):
|
||||
"""
|
||||
Retrieve a value from the global state using the given key.
|
||||
"""
|
||||
if not key or use_default:
|
||||
return (default_value,)
|
||||
|
||||
# Get the shared state dictionary
|
||||
shared_state = self.state_manager.get_state("__shared_keys__", {})
|
||||
|
||||
# Check if the key exists
|
||||
if key in shared_state:
|
||||
return (shared_state[key],)
|
||||
|
||||
# Return default value if key not found
|
||||
return (default_value,)
|
||||
|
||||
class SetStateNode(ControlNodeBase):
|
||||
"""
|
||||
Node that stores a value in the global state with a user-specified key.
|
||||
The value will be accessible in future workflow runs through GetStateNode.
|
||||
"""
|
||||
CATEGORY = "utils"
|
||||
RETURN_TYPES = (AlwaysEqualProxy("*"),)
|
||||
RETURN_NAMES = ("value",)
|
||||
FUNCTION = "update"
|
||||
OUTPUT_NODE = True
|
||||
DESCRIPTION = "Store a value in the global state with the given key. The value will be accessible in future workflow runs through GetStateNode."
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
inputs = super().INPUT_TYPES()
|
||||
inputs["required"].update({
|
||||
"key": ("STRING", {"default": "default_key", "tooltip": "The key to store the value under. If not provided, the value will not be stored."}),
|
||||
"value": (AlwaysEqualProxy("*"), {"tooltip": "The value to store in the global state."}),
|
||||
})
|
||||
return inputs
|
||||
|
||||
def update(self, key: str, value, always_execute=True):
|
||||
"""
|
||||
Store a value in the global state with the given key.
|
||||
"""
|
||||
if not key:
|
||||
return (value,)
|
||||
|
||||
try:
|
||||
shared_state = self.state_manager.get_state("__shared_keys__", {})
|
||||
shared_state[key] = copy.deepcopy(value)
|
||||
self.state_manager.set_state("__shared_keys__", shared_state)
|
||||
except Exception as e:
|
||||
print(f"[State Node] Error storing value: {str(e)}")
|
||||
|
||||
return (value,)
|
||||
|
||||
|
||||
@@ -42,7 +42,7 @@ class HandTrackingNode:
|
||||
RETURN_NAMES = ("debug_image", "hand_data")
|
||||
FUNCTION = "track_hands"
|
||||
CATEGORY = "image/processing"
|
||||
|
||||
DESCRIPTION = "(((EXPERIMENTAL))) Track hands in an image using MediaPipe. Returns a debug image with landmarks and a dictionary of hand data."
|
||||
def __init__(self):
|
||||
self.last_detection_conf = 0.5
|
||||
self.last_tracking_conf = 0.5
|
||||
|
||||
@@ -3,6 +3,7 @@ import comfy.utils
|
||||
import numpy as np
|
||||
import base64
|
||||
import re
|
||||
import math
|
||||
from io import BytesIO
|
||||
from PIL import Image, ImageFont, ImageDraw, ImageOps
|
||||
import nodes
|
||||
@@ -574,4 +575,439 @@ class LoadImageFromPath_:
|
||||
mask = mask.unsqueeze(0)
|
||||
|
||||
return (image, mask)
|
||||
|
||||
class FeedbackEffect:
|
||||
"""
|
||||
Ultra-fast real-time visual feedback processor designed for SetState/GetState loops.
|
||||
Creates stunning, hypnotic visuals with optimized GPU-accelerated transformations.
|
||||
|
||||
For best results: connect to a SetState/GetState loop and adjust parameters in real-time.
|
||||
"""
|
||||
CATEGORY = "image/effects"
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("image",)
|
||||
FUNCTION = "apply_effect"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE", {"tooltip": "Input image (typically from GetStateNode)"}),
|
||||
"effect": (["fractal_zoom", "kaleidoscope", "plasma_field", "quantum_wave",
|
||||
"neural_dream", "digital_flow", "hypershock"], {
|
||||
"default": "fractal_zoom",
|
||||
"tooltip": "Type of real-time effect to apply"
|
||||
}),
|
||||
"intensity": ("FLOAT", {
|
||||
"default": 0.65,
|
||||
"min": 0.1,
|
||||
"max": 1.0,
|
||||
"step": 0.05,
|
||||
"tooltip": "Effect strength/intensity"
|
||||
}),
|
||||
"variation": ("INT", {
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"max": 5,
|
||||
"step": 1,
|
||||
"tooltip": "Variation/mode of the selected effect"
|
||||
}),
|
||||
"speed": ("FLOAT", {
|
||||
"default": 0.5,
|
||||
"min": 0.05,
|
||||
"max": 1.0,
|
||||
"step": 0.05,
|
||||
"tooltip": "Animation speed factor"
|
||||
}),
|
||||
"feedback_blend": ("FLOAT", {
|
||||
"default": 0.95,
|
||||
"min": 0.5,
|
||||
"max": 0.999,
|
||||
"step": 0.005,
|
||||
"tooltip": "How much of previous frame to retain (higher = more trail effect)"
|
||||
}),
|
||||
"color_shift": ("FLOAT", {
|
||||
"default": 0.2,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.05,
|
||||
"tooltip": "Amount of color transformation to apply"
|
||||
}),
|
||||
"reset": ("BOOLEAN", {
|
||||
"default": False,
|
||||
"tooltip": "Reset the feedback loop"
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
def __init__(self):
|
||||
self.frame_count = 0
|
||||
self.time = 0
|
||||
self.last_image = None
|
||||
self.last_effect = None
|
||||
self.effect_buffers = {}
|
||||
|
||||
def apply_effect(self, image, effect, intensity, variation, speed, feedback_blend, color_shift, reset):
|
||||
batch_size, height, width, channels = image.shape
|
||||
device = image.device
|
||||
|
||||
# Reset state if requested or if effect changes
|
||||
if reset or effect != self.last_effect:
|
||||
self.frame_count = 0
|
||||
self.time = 0
|
||||
self.effect_buffers = {}
|
||||
|
||||
# Update time and frame counter
|
||||
self.frame_count += 1
|
||||
frame_time = self.frame_count * speed * 0.03
|
||||
self.time = frame_time
|
||||
self.last_effect = effect
|
||||
|
||||
# Fast path for batch processing - we'll process each image separately
|
||||
result_images = []
|
||||
|
||||
for b in range(batch_size):
|
||||
img = image[b]
|
||||
result = self._process_single_image(
|
||||
img, effect, intensity, variation,
|
||||
feedback_blend, color_shift, device,
|
||||
width, height, frame_time, speed
|
||||
)
|
||||
result_images.append(result)
|
||||
|
||||
return (torch.stack(result_images),)
|
||||
|
||||
def _process_single_image(self, img, effect, intensity, variation,
|
||||
feedback_blend, color_shift, device,
|
||||
width, height, time, speed):
|
||||
# Highly optimized processing path - calculate grid transformations in one go
|
||||
# Coordinates grid (only calculate once and reuse)
|
||||
if not hasattr(self, 'grid_cache') or self.grid_cache['size'] != (height, width):
|
||||
y_norm, x_norm = torch.meshgrid(
|
||||
torch.linspace(-1, 1, height, device=device),
|
||||
torch.linspace(-1, 1, width, device=device),
|
||||
indexing='ij'
|
||||
)
|
||||
self.grid_cache = {
|
||||
'size': (height, width),
|
||||
'x_norm': x_norm,
|
||||
'y_norm': y_norm,
|
||||
'r': torch.sqrt(x_norm**2 + y_norm**2),
|
||||
'theta': torch.atan2(y_norm, x_norm)
|
||||
}
|
||||
|
||||
# Get cached grids
|
||||
x_norm = self.grid_cache['x_norm']
|
||||
y_norm = self.grid_cache['y_norm']
|
||||
r = self.grid_cache['r']
|
||||
theta = self.grid_cache['theta']
|
||||
|
||||
# Initialize source coordinates to identity
|
||||
source_x = x_norm.clone()
|
||||
source_y = y_norm.clone()
|
||||
|
||||
# Apply selected effect transformation
|
||||
if effect == "fractal_zoom":
|
||||
# Fractal zoom with organic movement
|
||||
zoom = 0.98 + math.sin(time * 2.5) * 0.02
|
||||
|
||||
# Variation affects spiral pattern
|
||||
spiral = variation * 0.2 * intensity
|
||||
angle = theta + r * spiral + time * speed
|
||||
|
||||
# Complex transform that creates evolving spiral patterns
|
||||
dx = torch.sin(r * (3 + variation) + time * 3) * 0.05 * intensity
|
||||
dy = torch.cos(r * (3 + variation) + time * 3.5) * 0.05 * intensity
|
||||
|
||||
# Apply zoom
|
||||
source_x = (x_norm * zoom + dx) * (1 + 0.03 * intensity * math.sin(time))
|
||||
source_y = (y_norm * zoom + dy) * (1 + 0.03 * intensity * math.cos(time))
|
||||
|
||||
# Add swirling effect
|
||||
twist = r * (4.0 * intensity) + time * (1.5 + variation * 0.3)
|
||||
source_x += torch.sin(twist + theta) * 0.05 * intensity
|
||||
source_y += torch.cos(twist - theta) * 0.05 * intensity
|
||||
|
||||
elif effect == "kaleidoscope":
|
||||
# Kaleidoscope effect with dynamic symmetry
|
||||
segments = 4 + variation * 2
|
||||
segment_angle = 2 * torch.pi / segments
|
||||
|
||||
# Rotate over time
|
||||
rot_theta = theta + time * (0.2 + 0.1 * variation)
|
||||
|
||||
# Mirror segments
|
||||
segment_id = torch.floor(rot_theta / segment_angle)
|
||||
in_segment = rot_theta - segment_id * segment_angle
|
||||
flip_mask = (segment_id % 2 == 0)
|
||||
mirror_theta = torch.where(flip_mask, in_segment, segment_angle - in_segment)
|
||||
|
||||
# Apply transform with zoom breathing and ripples
|
||||
zoom = 0.8 + 0.2 * math.sin(time * 1.5) + r * 0.1 * intensity
|
||||
ripple = torch.sin(r * (8 + variation) - time * 3) * 0.06 * intensity
|
||||
|
||||
source_x = torch.cos(mirror_theta) * r * (zoom + ripple)
|
||||
source_y = torch.sin(mirror_theta) * r * (zoom + ripple)
|
||||
|
||||
elif effect == "plasma_field":
|
||||
# Plasma with flow dynamics
|
||||
scale = 5.0 + variation * 2.0
|
||||
speed_factor = 3.0 + variation * 0.5
|
||||
|
||||
# Create plasma-like distortion fields
|
||||
plasma_x = torch.sin(x_norm * scale + time * speed_factor) * torch.cos(y_norm * (scale*0.8) + time)
|
||||
plasma_y = torch.sin(y_norm * scale - time * (speed_factor*0.7)) * torch.cos(x_norm * (scale*0.9) - time * 1.3)
|
||||
plasma_xy = torch.sin((x_norm + y_norm) * (scale*0.4) + time * 1.5) * 0.3
|
||||
|
||||
# Combine fields for rich movement
|
||||
flow_x = (plasma_x + plasma_xy) * 0.15 * intensity
|
||||
flow_y = (plasma_y + plasma_xy) * 0.15 * intensity
|
||||
|
||||
# Apply field distortion
|
||||
source_x = x_norm + flow_x
|
||||
source_y = y_norm + flow_y
|
||||
|
||||
elif effect == "quantum_wave":
|
||||
# Quantum wave effect with interference patterns
|
||||
time_factor = time * (2.0 + variation * 0.4)
|
||||
|
||||
# Dynamic wave parameters
|
||||
wave_scale = (4.0 + variation) + math.sin(time) * 1.5
|
||||
wave_speed = 1.5 + variation * 0.2
|
||||
|
||||
# Create interference pattern
|
||||
wave1 = torch.sin(x_norm * wave_scale + time_factor * wave_speed)
|
||||
wave2 = torch.sin(y_norm * wave_scale - time_factor * wave_speed)
|
||||
wave3 = torch.sin((x_norm + y_norm) * wave_scale * 0.5 + time_factor * 1.3)
|
||||
wave4 = torch.sin(r * wave_scale * 0.8 - time_factor)
|
||||
|
||||
# Combine wave patterns
|
||||
interference = (wave1 + wave2 + wave3 + wave4) * 0.25
|
||||
|
||||
# Apply distortion with intensity scaling
|
||||
dist_scale = 0.12 * intensity * (1.0 + 0.3 * math.sin(time * 0.5))
|
||||
source_x = x_norm + interference * dist_scale * (1.0 + 0.5 * x_norm)
|
||||
source_y = y_norm + interference * dist_scale * (1.0 + 0.5 * y_norm)
|
||||
|
||||
# Add subtle spiral
|
||||
rotation = time * 0.3 * (1 + variation * 0.1)
|
||||
rot_x = x_norm * torch.cos(torch.tensor(rotation, device=device)) - y_norm * torch.sin(torch.tensor(rotation, device=device))
|
||||
rot_y = x_norm * torch.sin(torch.tensor(rotation, device=device)) + y_norm * torch.cos(torch.tensor(rotation, device=device))
|
||||
|
||||
# Blend rotation into final coordinates
|
||||
source_x = source_x * 0.85 + rot_x * 0.15
|
||||
source_y = source_y * 0.85 + rot_y * 0.15
|
||||
|
||||
elif effect == "neural_dream":
|
||||
# Neural network-like patterns that evolve
|
||||
# Create layered activation patterns
|
||||
scale_factor = 3.0 + variation * 0.5
|
||||
time_offset = time * (1.2 + variation * 0.2)
|
||||
|
||||
# First activation layer - positional encodings
|
||||
act1 = torch.sin(x_norm * scale_factor + time_offset)
|
||||
act2 = torch.sin(y_norm * scale_factor * 1.1 - time_offset * 0.7)
|
||||
act3 = torch.sin((x_norm*0.5 + y_norm*1.5) * scale_factor * 0.7 + time_offset * 0.5)
|
||||
|
||||
# Second layer - non-linear combinations
|
||||
layer2_1 = torch.tanh(act1 * act2 * 0.5 + act3 * 0.3)
|
||||
layer2_2 = torch.tanh(act2 * act3 * 0.5 + act1 * 0.3)
|
||||
|
||||
# Output layer with residual connection
|
||||
dx = layer2_1 * 0.15 * intensity
|
||||
dy = layer2_2 * 0.15 * intensity
|
||||
|
||||
# Apply distortion with subtle spiral
|
||||
angle = r * (variation * 0.2) + time * (0.2 + variation * 0.05)
|
||||
source_x = x_norm + dx + torch.sin(angle) * r * 0.02 * intensity
|
||||
source_y = y_norm + dy + torch.cos(angle) * r * 0.02 * intensity
|
||||
|
||||
elif effect == "digital_flow":
|
||||
# Digital fluid-like motion with emergent patterns
|
||||
# Base flow field
|
||||
scale = 4.0 + variation
|
||||
time_factor = time * (1.0 + variation * 0.2)
|
||||
|
||||
# Create flow vectors using simplex-like patterns
|
||||
flow_x1 = torch.sin(x_norm * scale + time_factor) * torch.cos(y_norm * (scale*0.5) + time_factor * 0.6)
|
||||
flow_y1 = torch.cos(x_norm * (scale*0.6) - time_factor * 0.7) * torch.sin(y_norm * scale - time_factor)
|
||||
|
||||
# Second flow field at different scale
|
||||
flow_x2 = torch.sin((x_norm + y_norm) * (scale*0.4) + time_factor * 1.3)
|
||||
flow_y2 = torch.cos((y_norm - x_norm) * (scale*0.4) - time_factor * 1.1)
|
||||
|
||||
# Combine flow fields
|
||||
flow_x = (flow_x1 * 0.6 + flow_x2 * 0.4) * 0.12 * intensity
|
||||
flow_y = (flow_y1 * 0.6 + flow_y2 * 0.4) * 0.12 * intensity
|
||||
|
||||
# Apply flow with subtle zoom pulse
|
||||
zoom = 1.0 + math.sin(time * 1.5) * 0.01 * intensity
|
||||
source_x = x_norm * zoom + flow_x
|
||||
source_y = y_norm * zoom + flow_y
|
||||
|
||||
# Add digital quantization if higher variation
|
||||
if variation >= 3:
|
||||
# Quantize stronger for higher variations
|
||||
quant_level = 15 + (5 - variation) * 5 # More divisions for lower variation
|
||||
source_x = torch.round(source_x * quant_level) / quant_level
|
||||
source_y = torch.round(source_y * quant_level) / quant_level
|
||||
|
||||
elif effect == "hypershock":
|
||||
# High energy effect with shockwaves and energy fields
|
||||
# Base frequencies
|
||||
freq = 5.0 + variation
|
||||
time_factor = time * (2.0 + variation * 0.3)
|
||||
|
||||
# Create pulsing shockwaves
|
||||
shock1_radius = (time_factor * 0.7) % 2.0 - 1.0
|
||||
shock1 = torch.exp(-((r - shock1_radius) * 10)**2) * 0.5
|
||||
|
||||
shock2_radius = (time_factor * 0.5 + 1.0) % 2.0 - 1.0
|
||||
shock2 = torch.exp(-((r - shock2_radius) * 10)**2) * 0.4
|
||||
|
||||
# Energy field
|
||||
energy_x = torch.sin(x_norm * freq + time_factor * 1.5) * 0.1
|
||||
energy_y = torch.sin(y_norm * freq - time_factor * 1.3) * 0.1
|
||||
|
||||
# Combine effects
|
||||
dx = (shock1 + shock2 * 0.7) * x_norm * 0.3 + energy_x
|
||||
dy = (shock1 + shock2 * 0.7) * y_norm * 0.3 + energy_y
|
||||
|
||||
# Apply distortion with high energy scaling
|
||||
source_x = x_norm + dx * intensity
|
||||
source_y = y_norm + dy * intensity
|
||||
|
||||
# Add spiral warp for higher variations
|
||||
if variation >= 2:
|
||||
warp = torch.sin(r * (8 - variation) + time_factor) * 0.1 * intensity
|
||||
source_x += torch.sin(theta + torch.tensor(time, device=device)) * warp
|
||||
source_y += torch.cos(theta + torch.tensor(time, device=device)) * warp
|
||||
|
||||
# Ensure coordinates stay within bounds
|
||||
source_x = torch.clamp(source_x, -1.98, 1.98)
|
||||
source_y = torch.clamp(source_y, -1.98, 1.98)
|
||||
|
||||
# Convert from normalized coords to actual pixel sampling grid [-1, 1]
|
||||
grid = torch.stack([source_x, source_y], dim=-1)
|
||||
grid = grid.unsqueeze(0) # Add batch dimension
|
||||
|
||||
# Sample the image (fast tensor operation)
|
||||
img_rgb = img[..., :3].permute(2, 0, 1).unsqueeze(0) # [1, C, H, W]
|
||||
sampled = torch.nn.functional.grid_sample(
|
||||
img_rgb, grid,
|
||||
mode='bilinear',
|
||||
padding_mode='reflection',
|
||||
align_corners=False
|
||||
).squeeze(0).permute(1, 2, 0) # Back to [H, W, C]
|
||||
|
||||
# Apply color transformations based on selected effect
|
||||
if color_shift > 0:
|
||||
# Optimize color processing by pushing calculations to the GPU
|
||||
if effect == "fractal_zoom":
|
||||
# Vibrant color amplification
|
||||
color_factor = time * 0.1 % 1.0
|
||||
hsv_rotation = torch.tensor([
|
||||
[math.cos(color_factor * 6.28), -math.sin(color_factor * 6.28), 0],
|
||||
[math.sin(color_factor * 6.28), math.cos(color_factor * 6.28), 0],
|
||||
[0, 0, 1]
|
||||
], device=device)
|
||||
|
||||
flat_rgb = sampled.reshape(-1, 3)
|
||||
rotated_rgb = torch.matmul(flat_rgb, hsv_rotation)
|
||||
sampled = rotated_rgb.reshape(sampled.shape)
|
||||
|
||||
# Vibrance enhancement
|
||||
luma = sampled[..., 0] * 0.299 + sampled[..., 1] * 0.587 + sampled[..., 2] * 0.114
|
||||
chroma = sampled - luma.unsqueeze(-1)
|
||||
sampled = luma.unsqueeze(-1) + chroma * (1 + color_shift * 0.5)
|
||||
|
||||
elif effect in ["kaleidoscope", "hypershock"]:
|
||||
# High contrast with glow
|
||||
# Add bloom/glow to bright areas
|
||||
bloom_threshold = 0.7
|
||||
bloom_mask = (sampled.mean(dim=-1) > bloom_threshold).float().unsqueeze(-1)
|
||||
bloom = sampled * bloom_mask * color_shift * 0.5
|
||||
|
||||
# Apply bloom with additive blending
|
||||
sampled = torch.clamp(sampled + bloom, 0, 1)
|
||||
|
||||
# Edge enhancement for kaleidoscope
|
||||
if effect == "kaleidoscope":
|
||||
edge_x = torch.abs(sampled[:, 1:] - sampled[:, :-1]).mean(-1, keepdim=True)
|
||||
edge_x = torch.cat([edge_x, edge_x[:, -1:]], dim=1)
|
||||
|
||||
edge_y = torch.abs(sampled[1:, :] - sampled[:-1, :]).mean(-1, keepdim=True)
|
||||
edge_y = torch.cat([edge_y, edge_y[-1:, :]], dim=0)
|
||||
|
||||
edge = (edge_x + edge_y) * 0.5
|
||||
sampled = torch.clamp(sampled + edge * color_shift, 0, 1)
|
||||
|
||||
elif effect in ["plasma_field", "neural_dream"]:
|
||||
# Hue rotation with time
|
||||
hue_shift = (time * 0.05) % 1.0
|
||||
|
||||
# RGB to HLS-like transform (simplified)
|
||||
cmax, _ = torch.max(sampled, dim=-1, keepdim=True)
|
||||
cmin, _ = torch.min(sampled, dim=-1, keepdim=True)
|
||||
delta = cmax - cmin
|
||||
|
||||
# Calculate luminance
|
||||
luma = (cmax + cmin) * 0.5
|
||||
|
||||
# Calculate saturation
|
||||
sat = torch.where(delta > 0, delta / (1 - torch.abs(2 * luma - 1) + 1e-6), torch.zeros_like(delta))
|
||||
|
||||
# Enhance saturation
|
||||
sat = torch.clamp(sat * (1 + color_shift), 0, 1)
|
||||
|
||||
# Apply new saturation while preserving luminance
|
||||
normalized = (sampled - cmin) / (delta + 1e-6)
|
||||
shifted = luma + (normalized - 0.5) * sat * (1 - torch.abs(2 * luma - 1))
|
||||
sampled = torch.clamp(shifted, 0, 1)
|
||||
|
||||
elif effect == "quantum_wave":
|
||||
# Quantum interference colors
|
||||
# Shift RGB channels with interference pattern
|
||||
shift = torch.sin(r * 4 + time * 2) * 0.1 * color_shift
|
||||
|
||||
# RGB channel shifting
|
||||
r_shift = torch.roll(sampled[..., 0], shifts=int(shift[0,0].item() * width), dims=0)
|
||||
r_shift = torch.roll(r_shift, shifts=int(shift[0,0].item() * height), dims=1)
|
||||
|
||||
b_shift = torch.roll(sampled[..., 2], shifts=-int(shift[0,0].item() * width), dims=0)
|
||||
b_shift = torch.roll(b_shift, shifts=-int(shift[0,0].item() * height), dims=1)
|
||||
|
||||
sampled = torch.stack([r_shift, sampled[..., 1], b_shift], dim=-1)
|
||||
|
||||
elif effect == "digital_flow":
|
||||
# Digital color quantization
|
||||
if color_shift > 0.3:
|
||||
quant_levels = max(3, int(10 * (1 - color_shift))) # Fewer levels with higher color_shift
|
||||
sampled = torch.floor(sampled * quant_levels) / quant_levels
|
||||
|
||||
# Alpha handling
|
||||
if img.shape[-1] == 4:
|
||||
# Keep original alpha channel
|
||||
alpha = torch.nn.functional.grid_sample(
|
||||
img[..., 3:].permute(2, 0, 1).unsqueeze(0),
|
||||
grid,
|
||||
mode='bilinear',
|
||||
padding_mode='reflection',
|
||||
align_corners=False
|
||||
).squeeze(0).permute(1, 2, 0)
|
||||
|
||||
# Recombine with RGB
|
||||
sampled = torch.cat([sampled, alpha], dim=-1)
|
||||
|
||||
# Apply feedback blend if we have a last image
|
||||
if self.last_image is not None and img.shape == self.last_image.shape:
|
||||
# Optimize blend operation
|
||||
result = sampled * (1.0 - feedback_blend) + self.last_image * feedback_blend
|
||||
else:
|
||||
result = sampled
|
||||
|
||||
# Store for next frame
|
||||
self.last_image = result
|
||||
|
||||
return result
|
||||
|
||||
|
||||
@@ -0,0 +1,409 @@
|
||||
import { app } from "../../../scripts/app.js";
|
||||
|
||||
// Track all keys available
|
||||
const stateKeys = new Set();
|
||||
|
||||
// Track connections to visualize data flow between set and get nodes
|
||||
const connectionVisualizer = {
|
||||
// Map of key -> {setNode, getNodes}
|
||||
connections: {},
|
||||
|
||||
registerSetNode: function(node, key) {
|
||||
if (!key) return;
|
||||
|
||||
// Initialize if not exists
|
||||
if (!this.connections[key]) {
|
||||
this.connections[key] = { setNode: null, getNodes: [] };
|
||||
}
|
||||
|
||||
this.connections[key].setNode = node;
|
||||
|
||||
// Find corresponding get nodes
|
||||
this.updateConnections(key);
|
||||
},
|
||||
|
||||
registerGetNode: function(node, key) {
|
||||
if (!key) return;
|
||||
|
||||
// Initialize if not exists
|
||||
if (!this.connections[key]) {
|
||||
this.connections[key] = { setNode: null, getNodes: [] };
|
||||
}
|
||||
|
||||
// Add to getNodes if not already there
|
||||
if (!this.connections[key].getNodes.includes(node)) {
|
||||
this.connections[key].getNodes.push(node);
|
||||
}
|
||||
|
||||
// Update connections
|
||||
this.updateConnections(key);
|
||||
},
|
||||
|
||||
unregisterNode: function(node, key) {
|
||||
if (!key || !this.connections[key]) return;
|
||||
|
||||
// Remove from connections
|
||||
if (this.connections[key].setNode === node) {
|
||||
this.connections[key].setNode = null;
|
||||
}
|
||||
|
||||
const getNodes = this.connections[key].getNodes;
|
||||
const index = getNodes.indexOf(node);
|
||||
if (index !== -1) {
|
||||
getNodes.splice(index, 1);
|
||||
}
|
||||
|
||||
// Remove key if no nodes left
|
||||
if (!this.connections[key].setNode && this.connections[key].getNodes.length === 0) {
|
||||
delete this.connections[key];
|
||||
}
|
||||
},
|
||||
|
||||
updateConnections: function(key) {
|
||||
if (!key || !this.connections[key]) return;
|
||||
|
||||
const connection = this.connections[key];
|
||||
|
||||
// Update visual properties for set node
|
||||
if (connection.setNode) {
|
||||
connection.setNode.title = `Set State: ${key}`;
|
||||
}
|
||||
|
||||
// Update visual properties for get nodes
|
||||
connection.getNodes.forEach(getNode => {
|
||||
getNode.title = `Get State: ${key}`;
|
||||
});
|
||||
},
|
||||
|
||||
getSetNodeForKey: function(key) {
|
||||
return this.connections[key]?.setNode;
|
||||
},
|
||||
|
||||
getNodesForKey: function(key) {
|
||||
return this.connections[key]?.getNodes || [];
|
||||
}
|
||||
};
|
||||
|
||||
// Helper to show alerts to the user
|
||||
function showAlert(message) {
|
||||
app.extensionManager.toast.add({
|
||||
severity: 'info',
|
||||
summary: "State Nodes",
|
||||
detail: `${message}`,
|
||||
life: 3000,
|
||||
});
|
||||
}
|
||||
|
||||
// Helper to find all set nodes with a given key
|
||||
function findSetNodesWithKey(graph, key) {
|
||||
return graph._nodes.filter(node =>
|
||||
node.type === "SetStateNode" &&
|
||||
node.widgets &&
|
||||
node.widgets.find(w => w.name === "key" && w.value === key)
|
||||
);
|
||||
}
|
||||
|
||||
// Helper to find all get nodes with a given key
|
||||
function findGetNodesWithKey(graph, key) {
|
||||
return graph._nodes.filter(node =>
|
||||
node.type === "GetStateNode" &&
|
||||
node.widgets &&
|
||||
node.widgets.find(w => w.name === "key" && w.value === key)
|
||||
);
|
||||
}
|
||||
|
||||
// Helper to find all get nodes in the graph
|
||||
function findAllGetNodes(graph) {
|
||||
return graph._nodes.filter(node => node.type === "GetStateNode");
|
||||
}
|
||||
|
||||
// Helper to update all get nodes with the current list of keys
|
||||
function updateAllGetNodesDropdowns(graph) {
|
||||
const getNodes = findAllGetNodes(graph);
|
||||
getNodes.forEach(getNode => {
|
||||
const keyWidget = getNode.widgets.find(w => w.name === "key");
|
||||
if (keyWidget && keyWidget.options) {
|
||||
// Update the values in the dropdown
|
||||
keyWidget.options.values = [...stateKeys].sort();
|
||||
|
||||
// If the current key no longer exists, reset it
|
||||
if (keyWidget.value && !stateKeys.has(keyWidget.value)) {
|
||||
keyWidget.value = "";
|
||||
getNode.title = "Get State";
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// Register the SetStateNode for UI customization
|
||||
app.registerExtension({
|
||||
name: "SetStateNode",
|
||||
async beforeRegisterNodeDef(nodeType, nodeData) {
|
||||
if (nodeData.name !== "SetStateNode") return;
|
||||
|
||||
// Add custom widgets and behavior
|
||||
const onNodeCreated = nodeType.prototype.onNodeCreated;
|
||||
nodeType.prototype.onNodeCreated = function() {
|
||||
if (onNodeCreated) {
|
||||
onNodeCreated.apply(this, arguments);
|
||||
}
|
||||
|
||||
// Store reference to the node
|
||||
const node = this;
|
||||
|
||||
// Track key changes
|
||||
const keyWidget = node.widgets.find(w => w.name === "key");
|
||||
if (keyWidget) {
|
||||
const originalCallback = keyWidget.callback;
|
||||
keyWidget.callback = function(value) {
|
||||
if (originalCallback) {
|
||||
originalCallback.apply(this, arguments);
|
||||
}
|
||||
|
||||
// Add to available keys
|
||||
if (value && value.trim() !== "") {
|
||||
stateKeys.add(value);
|
||||
node.title = "Set State: " + value;
|
||||
|
||||
// Update all GetStateNodes
|
||||
if (node.graph) {
|
||||
updateAllGetNodesDropdowns(node.graph);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Initialize if there's a default value
|
||||
if (keyWidget.value && keyWidget.value.trim() !== "") {
|
||||
stateKeys.add(keyWidget.value);
|
||||
node.title = "Set State: " + keyWidget.value;
|
||||
}
|
||||
}
|
||||
|
||||
// Set node appearance
|
||||
this.color = "#2a363b";
|
||||
this.bgcolor = "#3f5159";
|
||||
|
||||
// Add extra context menu items
|
||||
const getExtraMenuOptions = node.getExtraMenuOptions;
|
||||
node.getExtraMenuOptions = function(_, options) {
|
||||
if (getExtraMenuOptions) {
|
||||
getExtraMenuOptions.apply(this, arguments);
|
||||
}
|
||||
|
||||
const keyWidget = node.widgets.find(w => w.name === "key");
|
||||
if (keyWidget && keyWidget.value) {
|
||||
const key = keyWidget.value;
|
||||
|
||||
// Add menu to see all connected get nodes
|
||||
const getNodes = findGetNodesWithKey(this.graph, key);
|
||||
if (getNodes.length > 0) {
|
||||
const getNodesSubmenu = getNodes.map(getNode => ({
|
||||
content: `Get Node: ${getNode.id}`,
|
||||
callback: () => {
|
||||
app.canvas.centerOnNode(getNode);
|
||||
app.canvas.selectNode(getNode);
|
||||
}
|
||||
}));
|
||||
|
||||
options.unshift({
|
||||
content: "Connected Get Nodes",
|
||||
has_submenu: true,
|
||||
submenu: {
|
||||
options: getNodesSubmenu
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// When the node is removed
|
||||
const onRemoved = node.onRemoved;
|
||||
node.onRemoved = function() {
|
||||
if (onRemoved) {
|
||||
onRemoved.apply(this, arguments);
|
||||
}
|
||||
|
||||
// Check if we need to remove the key from the available keys
|
||||
const keyWidget = node.widgets.find(w => w.name === "key");
|
||||
if (keyWidget && keyWidget.value) {
|
||||
const key = keyWidget.value;
|
||||
// Only remove if no other SetStateNodes are using this key
|
||||
const otherSetNodes = findSetNodesWithKey(this.graph, key);
|
||||
if (otherSetNodes.length <= 1) { // 1 because this node is still in the list
|
||||
stateKeys.delete(key);
|
||||
|
||||
// Update the dropdowns in all GetStateNodes
|
||||
updateAllGetNodesDropdowns(this.graph);
|
||||
}
|
||||
}
|
||||
};
|
||||
};
|
||||
}
|
||||
});
|
||||
|
||||
// Register the GetStateNode for UI customization
|
||||
app.registerExtension({
|
||||
name: "GetStateNode",
|
||||
async beforeRegisterNodeDef(nodeType, nodeData) {
|
||||
if (nodeData.name !== "GetStateNode") return;
|
||||
|
||||
// Add custom widgets and behavior
|
||||
const onNodeCreated = nodeType.prototype.onNodeCreated;
|
||||
nodeType.prototype.onNodeCreated = function() {
|
||||
if (onNodeCreated) {
|
||||
onNodeCreated.apply(this, arguments);
|
||||
}
|
||||
|
||||
// Store reference to the node
|
||||
const node = this;
|
||||
|
||||
// Set node appearance
|
||||
this.color = "#2a363b";
|
||||
this.bgcolor = "#3f5159";
|
||||
|
||||
// Find the key widget and enhance it
|
||||
const keyWidget = node.widgets.find(w => w.name === "key");
|
||||
if (keyWidget) {
|
||||
// Replace widget with combo
|
||||
const index = node.widgets.indexOf(keyWidget);
|
||||
const keyValue = keyWidget.value; // Save current value
|
||||
|
||||
// Remove old widget
|
||||
node.widgets.splice(index, 1);
|
||||
|
||||
// Add new combo widget
|
||||
const newWidget = node.addWidget("combo", "key", keyValue, function(value) {
|
||||
if (value) {
|
||||
node.title = "Get State: " + value;
|
||||
} else {
|
||||
node.title = "Get State";
|
||||
}
|
||||
}, {
|
||||
values: () => [...stateKeys].sort(),
|
||||
name: "key", // Ensure name matches for server
|
||||
});
|
||||
|
||||
// Move it to the correct position
|
||||
if (index !== node.widgets.length - 1) {
|
||||
const removed = node.widgets.pop();
|
||||
node.widgets.splice(index, 0, removed);
|
||||
}
|
||||
|
||||
// Initialize title
|
||||
if (newWidget.value) {
|
||||
node.title = "Get State: " + newWidget.value;
|
||||
} else {
|
||||
node.title = "Get State";
|
||||
}
|
||||
}
|
||||
|
||||
// Add extra context menu items
|
||||
const getExtraMenuOptions = node.getExtraMenuOptions;
|
||||
node.getExtraMenuOptions = function(_, options) {
|
||||
if (getExtraMenuOptions) {
|
||||
getExtraMenuOptions.apply(this, arguments);
|
||||
}
|
||||
|
||||
const keyWidget = node.widgets.find(w => w.name === "key");
|
||||
if (keyWidget && keyWidget.value) {
|
||||
const key = keyWidget.value;
|
||||
|
||||
// Add menu to go to set nodes with this key
|
||||
const setNodes = findSetNodesWithKey(this.graph, key);
|
||||
if (setNodes.length > 0) {
|
||||
const setNodesSubmenu = setNodes.map(setNode => ({
|
||||
content: `Set Node: ${setNode.id}`,
|
||||
callback: () => {
|
||||
app.canvas.centerOnNode(setNode);
|
||||
app.canvas.selectNode(setNode);
|
||||
}
|
||||
}));
|
||||
|
||||
options.unshift({
|
||||
content: setNodes.length === 1 ? "Go to Set Node" : "Go to Set Nodes",
|
||||
has_submenu: setNodes.length > 1,
|
||||
callback: setNodes.length === 1 ? (() => {
|
||||
app.canvas.centerOnNode(setNodes[0]);
|
||||
app.canvas.selectNode(setNodes[0]);
|
||||
}) : undefined,
|
||||
submenu: setNodes.length > 1 ? {
|
||||
options: setNodesSubmenu
|
||||
} : undefined
|
||||
});
|
||||
}
|
||||
|
||||
// Add menu to see other get nodes with the same key
|
||||
const getNodes = findGetNodesWithKey(this.graph, key).filter(n => n !== this);
|
||||
if (getNodes.length > 0) {
|
||||
const getNodesSubmenu = getNodes.map(getNode => ({
|
||||
content: `Get Node: ${getNode.id}`,
|
||||
callback: () => {
|
||||
app.canvas.centerOnNode(getNode);
|
||||
app.canvas.selectNode(getNode);
|
||||
}
|
||||
}));
|
||||
|
||||
options.unshift({
|
||||
content: "Other Get Nodes",
|
||||
has_submenu: true,
|
||||
submenu: {
|
||||
options: getNodesSubmenu
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
};
|
||||
};
|
||||
}
|
||||
});
|
||||
|
||||
// Register the other state-related nodes
|
||||
app.registerExtension({
|
||||
name: "StateUtilityNodes",
|
||||
async beforeRegisterNodeDef(nodeType, nodeData) {
|
||||
if (nodeData.name !== "ClearStateNode" && nodeData.name !== "ListStatesNode") return;
|
||||
|
||||
// Add custom appearance
|
||||
const onNodeCreated = nodeType.prototype.onNodeCreated;
|
||||
nodeType.prototype.onNodeCreated = function() {
|
||||
if (onNodeCreated) {
|
||||
onNodeCreated.apply(this, arguments);
|
||||
}
|
||||
|
||||
// Set node appearance
|
||||
this.color = "#2a363b";
|
||||
this.bgcolor = "#3f5159";
|
||||
|
||||
// When ClearStateNode is executed, also update the UI
|
||||
if (nodeData.name === "ClearStateNode") {
|
||||
const onExecuted = this.onExecuted;
|
||||
this.onExecuted = function(message) {
|
||||
if (onExecuted) {
|
||||
onExecuted.apply(this, arguments);
|
||||
}
|
||||
|
||||
// Check if all states were cleared
|
||||
const action = this.widgets.find(w => w.name === "action");
|
||||
if (action && action.value === "Clear All States") {
|
||||
// Clear the state keys set
|
||||
stateKeys.clear();
|
||||
showAlert("All states cleared");
|
||||
|
||||
// Update all GetStateNodes
|
||||
updateAllGetNodesDropdowns(this.graph);
|
||||
} else {
|
||||
// Clear just one key
|
||||
const keyWidget = this.widgets.find(w => w.name === "key");
|
||||
if (keyWidget && keyWidget.value) {
|
||||
stateKeys.delete(keyWidget.value);
|
||||
showAlert(`State '${keyWidget.value}' cleared`);
|
||||
|
||||
// Update all GetStateNodes
|
||||
updateAllGetNodesDropdowns(this.graph);
|
||||
}
|
||||
}
|
||||
};
|
||||
}
|
||||
};
|
||||
}
|
||||
});
|
||||
Reference in New Issue
Block a user