Merge pull request #8 from ryanontheinside/refactor/initial-refactor

Refactor/initial refactor

Author: ryanontheinside

__init__.py

@@ -1,59 +1,82 @@
-from .controls.value_controls import FloatControl, IntControl, StringControl
-from .controls.sequence_controls import FloatSequence, IntSequence, StringSequence
-from .controls.utility_controls import FPSMonitor, SimilarityFilter, LazyCondition, LogicOperator, IfThenElse   
-from .controls.state_management_controls import StateResetNode, StateTestNode
-from .controls.motion_controls import MotionController, ROINode, IntegerMotionController
-from .misc_nodes import( 
-    DTypeConverter,
-    FastWebcamCapture,
-    YOLOSimilarityCompare, 
-    TextRenderer,
-    QuickShapeMask,
-    MultilineText, 
-    LoadImageFromPath_
-)
-from .media_pipe_nodes import HandTrackingNode, HandMaskNode
-from .controls.mask_controls import RepulsiveMaskNode, ResizeMaskNode
-
+import os
+import importlib
+import inspect
 import re
+import traceback
 
+# Directory containing node wrapper modules relative to this __init__.py
+WRAPPER_DIR = "node_wrappers.realtimenodes"
+NODE_CLASS_MAPPINGS = {}
+NODE_DISPLAY_NAME_MAPPINGS = {}
 
-NODE_CLASS_MAPPINGS = {
-    "FloatControl": FloatControl,
-    "IntControl": IntControl,
-    "StringControl": StringControl,
-    "FloatSequence": FloatSequence,
-    "IntSequence": IntSequence,
-    "StringSequence": StringSequence,
-    "FPSMonitor": FPSMonitor,
-    "SimilarityFilter": SimilarityFilter,
-    "LazyCondition": LazyCondition,
-    "LogicOperator": LogicOperator,
-    "IfThenElse": IfThenElse,
-    "MotionController": MotionController,
-    "IntegerMotionController": IntegerMotionController,
-    "YOLOSimilarityCompare": YOLOSimilarityCompare,
-    "TextRenderer": TextRenderer,
-    "ROINode": ROINode,
-    "QuickShapeMask": QuickShapeMask,
-    "DTypeConverter": DTypeConverter,
-    "FastWebcamCapture": FastWebcamCapture,
-    "MultilineText": MultilineText,
-    "LoadImageFromPath_": LoadImageFromPath_,
-    "HandTrackingNode": HandTrackingNode,
-    "HandMaskNode": HandMaskNode,
-    "RepulsiveMaskNode": RepulsiveMaskNode,
-    "ResizeMaskNode": ResizeMaskNode,
-    "StateResetNode": StateResetNode,
-    "StateTestNode": StateTestNode,
-}
+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_display_names = {}
+    
+    # Calculate absolute path to the package directory
+    package_abs_path = os.path.join(os.path.dirname(__file__), *package_path.split('.'))
+    
+    if not os.path.isdir(package_abs_path):
+        print(f"\033[91m" + f"[ComfyUI_RealtimeNodes] Warning: Directory not found: {package_abs_path}")
+        return loaded_nodes, loaded_display_names
 
+    print(f"\033[92m" + f"[ComfyUI_RealtimeNodes] Searching for nodes in: {package_abs_path}")
 
+    for filename in os.listdir(package_abs_path):
+        if filename.endswith(".py") and not filename.startswith('__'):
+            module_name = filename[:-3]
+            full_module_path = f".{package_path}.{module_name}" # Relative import path
+            
+            try:
+                module = importlib.import_module(full_module_path, package=__name__)
+                
+                # Find classes defined in the module
+                for name, obj in inspect.getmembers(module):
+                    if inspect.isclass(obj):
+                        # Check if it's a ComfyUI node (heuristic: has INPUT_TYPES)
+                        if hasattr(obj, 'INPUT_TYPES') and callable(obj.INPUT_TYPES):
+                             # Exclude base classes/mixins if they are explicitly named or lack a CATEGORY
+                             if not name.endswith("Base") and not name.endswith("Mixin") and hasattr(obj, 'CATEGORY'):
+                                loaded_nodes[name] = obj
+                                
+                                # Generate display name (similar to original logic)
+                                display_name = ' '.join(word.capitalize() for word in re.findall(r'[A-Z]?[a-z]+|[A-Z]+(?=[A-Z][a-z]|\d|\W|$)|\d+', name))
+                                suffix = " 🕒🅡🅣🅝"
+                                if not display_name.endswith(suffix):
+                                    display_name += suffix
+                                loaded_display_names[name] = display_name
+                                print(f"\033[94m" + f"[ComfyUI_RealtimeNodes]   - Loaded node: {name} -> {display_name}")
+                                
+            except ImportError as e:
+                print(f"\033[91m" + f"[ComfyUI_RealtimeNodes] Error importing module {full_module_path}: {e}")
+                traceback.print_exc()
+            except Exception as e:
+                print(f"\033[91m" + f"[ComfyUI_RealtimeNodes] Error processing module {full_module_path}: {e}")
+                traceback.print_exc()
+                
+    return loaded_nodes, loaded_display_names
 
+# --- 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
+]
+
+# 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)
 
-NODE_DISPLAY_NAME_MAPPINGS = {}
 
 suffix = " 🕒🅡🅣🅝"
 
@@ -70,8 +93,9 @@
 
     # Assign the final display name to the mappings
     NODE_DISPLAY_NAME_MAPPINGS[node_name] = display_name
+# --- Original Web Directory and Export ---
+WEB_DIRECTORY = "./web/js" # Adjusted path if needed
 
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
 
-WEB_DIRECTORY = "./web/js"
-
-__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] 
+print("\033[92m" + f"[ComfyUI_RealtimeNodes] Loaded {len(NODE_CLASS_MAPPINGS)} nodes.") 

node_wrappers/__init__.py

@@ -0,0 +1 @@
+# This file intentionally left blank 

node_wrappers/realtimenodes/__init__.py

@@ -0,0 +1 @@
+# This file intentionally left blank 

node_wrappers/realtimenodes/controls/__init__.py

@@ -0,0 +1 @@
+# This file intentionally left blank 

controls/mask_controls.py renamed to node_wrappers/realtimenodes/controls/mask_controls.py

@@ -1,31 +1,8 @@
 import torch
 import numpy as np
 import cv2
-from ..base.control_base import ControlNodeBase
-
-class MaskControlMixin:
-    """Mixin providing common mask functionality"""
-    def create_circle_mask(self, height, width, center_y, center_x, size):
-        """Create a circular mask with anti-aliased edges"""
-        y, x = np.ogrid[:height, :width]
-        radius = size * min(height, width) / 2
-        dist = np.sqrt((x - center_x)**2 + (y - center_y)**2)
-        mask = np.clip(radius + 1 - dist, 0, 1).astype(np.float32)
-        return mask
-
-    def get_initial_state(self, x_pos, y_pos, size, min_size=None, max_size=None):
-        """Get initial state dictionary"""
-        state = {
-            "x": x_pos,
-            "y": y_pos,
-            "size": size,
-            "initialized": True
-        }
-        if min_size is not None:
-            state["min_size"] = min_size
-        if max_size is not None:
-            state["max_size"] = max_size
-        return state
+from ....src.realtimenodes.control_base import ControlNodeBase
+from ....src.realtimenodes.mask_controls import MaskControlMixin
 
 class RepulsiveMaskNode(ControlNodeBase, MaskControlMixin):
     """Node that maintains a mask that repulses or attracts to input masks"""

controls/motion_controls.py renamed to node_wrappers/realtimenodes/controls/motion_controls.py

@@ -1,4 +1,4 @@
-from ..base.control_base import ControlNodeBase
+from ....src.realtimenodes.control_base import ControlNodeBase
 from enum import Enum
 import numpy as np
 import torch

controls/sequence_controls.py renamed to node_wrappers/realtimenodes/controls/sequence_controls.py

@@ -1,4 +1,4 @@
-from ..base.control_base import ControlNodeBase
+from ....src.realtimenodes.control_base import ControlNodeBase
 
 class SequenceControlBase(ControlNodeBase):
     """Base class for sequence-based controls"""

controls/state_management_controls.py renamed to node_wrappers/realtimenodes/controls/state_management_controls.py

@@ -1,4 +1,4 @@
-from ..base.control_base import ControlNodeBase
+from ....src.realtimenodes.control_base import ControlNodeBase
 
 class StateResetNode(ControlNodeBase):
     """Node that resets all control node states when triggered"""

controls/value_controls.py renamed to node_wrappers/realtimenodes/controls/value_controls.py

@@ -1,5 +1,5 @@
-from ..base.control_base import ControlNodeBase
-from ..patterns.movement_patterns import MOVEMENT_PATTERNS
+from ....src.realtimenodes.control_base import ControlNodeBase
+from ....src.realtimenodes.patterns.movement_patterns import MOVEMENT_PATTERNS
 
 class ValueControlBase(ControlNodeBase):
     """Base class for float and integer control nodes"""

node_wrappers/realtimenodes/logic/logic.py

@@ -0,0 +1,107 @@
+from ....src.utils.general_utils import AlwaysEqualProxy
+from ....src.realtimenodes.control_base import ControlNodeBase
+
+class LazyCondition(ControlNodeBase):
+    DESCRIPTION = "Uses lazy evaluation to truly skip execution of unused paths. Maintains state of the last value to circumvent feedback loops."
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "condition": (AlwaysEqualProxy("*"), {
+                    "tooltip": "When truthy (non-zero, non-empty, non-None), evaluates and returns if_true path. When falsy, returns either fallback or previous state of if_true.",
+                    "forceInput": True,
+                }),
+                "if_true": (AlwaysEqualProxy("*"), {
+                    "lazy": True,
+                    "tooltip": "The path that should only be evaluated when condition is truthy"
+                }),
+                "fallback": (AlwaysEqualProxy("*"), {
+                    "tooltip": "Alternative value to use when condition is falsy or no previous state of if_true"
+                }),
+                "use_fallback": ("BOOLEAN", {
+                    "default": False,
+                    "tooltip": "When False, uses last successful if_true result (if one exists). When True, uses fallback value"
+                }),
+            }
+        }
+    
+    RETURN_TYPES = (AlwaysEqualProxy("*"),)
+    FUNCTION = "update"
+    CATEGORY = "real-time/control/utility"
+
+    def check_lazy_status(self, condition, if_true, fallback, use_fallback):
+        """Only evaluate the if_true path if condition is truthy."""
+        needed = ["fallback"]  # Always need the fallback value
+        if condition:
+            needed.append("if_true")
+        return needed
+
+    def update(self, condition, if_true, fallback, use_fallback):
+        """Route to either if_true output or fallback value, maintaining state of last if_true."""
+        state = self.get_state({
+            "prev_output": None
+        })
+
+        if condition:  # Let Python handle truthiness
+            # Update last state when we run if_true path
+            state["prev_output"] = if_true if if_true is not None else fallback
+            if hasattr(if_true, 'detach'):
+                state["prev_output"] = if_true.detach().clone()
+            self.set_state(state)
+            return (if_true,)
+        else:
+            if use_fallback or state["prev_output"] is None:
+                return (fallback,)
+            else:
+                return (state["prev_output"],)
+
+class LogicOperator(ControlNodeBase):
+    DESCRIPTION = "Performs logical operations (AND, OR, NOT, XOR) on inputs based on their truthiness"
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "operation": (["AND", "OR", "NOT", "XOR"], {
+                    "default": "AND",
+                    "tooltip": "Logical operation to perform"
+                }),
+                "input_a": (AlwaysEqualProxy("*"), {
+                    "tooltip": "First input to evaluate for truthiness",
+                    "forceInput": True,
+                }),
+                "always_execute": ("BOOLEAN", {
+                    "default": True,
+                }),
+            },
+            "optional": {
+                "input_b": (AlwaysEqualProxy("*"), {
+                    "tooltip": "Second input to evaluate for truthiness (not used for NOT operation)",
+                }),
+            }
+        }
+    
+    RETURN_TYPES = ("BOOLEAN",)
+    FUNCTION = "update"
+    CATEGORY = "real-time/control/logic"
+
+    def update(self, operation, input_a, always_execute=True, input_b=None):
+        """Perform the selected logical operation on inputs based on their truthiness."""
+        a = bool(input_a)
+        
+        if operation == "NOT":
+            return (not a,)
+        
+        # For all other operations, we need input_b
+        b = bool(input_b)
+        
+        if operation == "AND":
+            return (a and b,)
+        elif operation == "OR":
+            return (a or b,)
+        elif operation == "XOR":
+            return (a != b,)
+        
+        # Should never get here, but just in case
+        return (False,)

node_wrappers/realtimenodes/media/__init__.py

@@ -0,0 +1 @@
+# This file intentionally left blank 

media_pipe_nodes.py renamed to node_wrappers/realtimenodes/media/media_pipe_nodes.py

@@ -0,0 +1 @@
+# This file intentionally left blank 

node_wrappers/realtimenodes/specialized/__init__.py

@@ -0,0 +1 @@
+# This file intentionally left blank