feat: draft on nodes refractor

__init__.py

@@ -1,36 +1,78 @@
+"""Setup expose Comfystream UI and native nodes to ComfyUI."""
+
 import os
 import pathlib
+import logging
+import importlib
+import traceback
+
+logging.basicConfig(level=logging.INFO)
+log = logging.getLogger(__name__)
+
+
+def load_nodes() -> tuple[dict[str, type], dict[str, str]]:
+    """Load all nodes in the Stream Pack."""
+    nodes_dir = pathlib.Path(__file__).parent / "nodes"
+    node_class_mappings, node_display_name_mappings = {}, {}
+
+    # Dynamically import all Python modules in the nodes directory.
+    for module_path in nodes_dir.iterdir():
+        if (
+            module_path.is_file()
+            and module_path.suffix == ".py"
+            and module_path.stem != "__init__"
+        ):
+            try:
+                module_name = f"{__package__}.nodes.{module_path.stem}"
+                module = importlib.import_module(module_name)
+
+                # Update mappings if defined in the module.
+                if hasattr(module, "NODE_CLASS_MAPPINGS"):
+                    node_class_mappings.update(module.NODE_CLASS_MAPPINGS)
+                if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS"):
+                    node_display_name_mappings.update(module.NODE_DISPLAY_NAME_MAPPINGS)
+            except Exception:
+                format_exc = traceback.format_exc()
+                log.error(f"Failed to load module {module_name}:\n{format_exc}")
+
+    return node_class_mappings, node_display_name_mappings
+
 
 def ensure_init_files():
-    """Create __init__.py files in comfy/ and comfy_extras/ directories if they don't exist"""
-    # Go up two levels from custom_nodes/comfystream_inside to reach ComfyUI root
-    comfy_root = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
-    base_dirs = ['comfy', 'comfy_extras']
+    """Ensures that the directories 'comfy/' and 'comfy_extras/' and their subdirectories
+    are recognized by ComfyUI as Python packages by creating empty __init__.py files
+    where needed.
+    """
+    # Go up two levels from custom_nodes/comfystream_inside to reach ComfyUI root.
+    comfy_root = pathlib.Path(__file__).resolve().parents[2]
+    base_dirs = ["comfy", "comfy_extras"]
+
     for base_dir in base_dirs:
-        base_path = os.path.join(comfy_root, base_dir)
-        if not os.path.exists(base_path):
+        base_path = comfy_root / base_dir
+        if not base_path.exists():
             continue
-            
-        # Create __init__.py in the root of base_dir first
-        root_init = os.path.join(base_path, "__init__.py")
-        if not os.path.exists(root_init):
-            with open(root_init, 'w') as f:
-                f.write("")
-                
-        # Then walk subdirectories
-        for root, dirs, files in os.walk(base_path):
-            init_path = os.path.join(root, "__init__.py")
-            if not os.path.exists(init_path):
-                with open(init_path, 'w') as f:
-                    f.write("")
-
-# Create __init__.py files in ComfyUI directories
+
+        # Create __init__.py in the root of base_dir first.
+        root_init = base_path / "__init__.py"
+        if not root_init.exists():
+            root_init.touch()
+
+        # Then walk subdirectories.
+        for subdir in base_path.rglob("*"):
+            if subdir.is_dir():
+                init_path = subdir / "__init__.py"
+                if not init_path.exists():
+                    init_path.touch()
+
+
+# Create __init__.py files in ComfyUI directories.
 ensure_init_files()
 
-# Point to the directory containing our web files
-WEB_DIRECTORY = "./nodes/web/js"
+# Point to the directory containing our web files.
+os.path.join(os.path.dirname(os.path.realpath(__file__)), "web")
 
-# Import and expose node classes
-from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+# Import and expose native nodes.
+NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS = load_nodes()
+NODE_DISPLAY_NAME_MAPPINGS["Comfystream"] = "Comfystream Native Nodes"
 
-__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] 
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

nodes/__init__.py

@@ -1,26 +1,2 @@
-"""ComfyStream nodes package"""
-
-from .audio_utils import *
-from .tensor_utils import *
-from .video_stream_utils import *
-from .api import *
-from .web import *
-
-# Collect all NODE_CLASS_MAPPINGS and NODE_DISPLAY_NAME_MAPPINGS from submodules
-NODE_CLASS_MAPPINGS = {}
-NODE_DISPLAY_NAME_MAPPINGS = {}
-
-# Import and update mappings from submodules
-for module in [audio_utils, tensor_utils, video_stream_utils, api, web]:
-    if hasattr(module, 'NODE_CLASS_MAPPINGS'):
-        NODE_CLASS_MAPPINGS.update(module.NODE_CLASS_MAPPINGS)
-    if hasattr(module, 'NODE_DISPLAY_NAME_MAPPINGS'):
-        NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS)
-
-# Web directory for UI components
-import os
-WEB_DIRECTORY = os.path.join(os.path.dirname(os.path.realpath(__file__)), "web")
-
-NODE_DISPLAY_NAME_MAPPINGS["ComfyStreamLauncher"] = "Launch ComfyStream 🚀"
-
-__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
+"""Contains the ComfyUI nodes that are installed natively with the ComfyStream package.
+"""

nodes/audio_utils/__init__.py

@@ -2,6 +2,10 @@
 from .save_audio_tensor import SaveAudioTensor
 from .pitch_shift import PitchShifter
 
-NODE_CLASS_MAPPINGS = {"LoadAudioTensor": LoadAudioTensor, "SaveAudioTensor": SaveAudioTensor, "PitchShifter": PitchShifter}
+NODE_CLASS_MAPPINGS = {
+    "LoadAudioTensor": LoadAudioTensor,
+    "SaveAudioTensor": SaveAudioTensor,
+    "PitchShifter": PitchShifter,
+}
 
 __all__ = ["NODE_CLASS_MAPPINGS"]

nodes/video_stream_utils/__init__.py

@@ -1,6 +1,6 @@
 """Video stream utility nodes for ComfyStream"""
 
-from .primary_input_load_image import PrimaryInputLoadImage 
+from .primary_input_load_image import PrimaryInputLoadImage
 
 NODE_CLASS_MAPPINGS = {"PrimaryInputLoadImage": PrimaryInputLoadImage}
 NODE_DISPLAY_NAME_MAPPINGS = {}

nodes/web/__init__.py

@@ -1,6 +1,5 @@
 """ComfyStream Web UI nodes"""
-import os
-import folder_paths
+
 
 # Define a simple Python class for the UI Preview node
 class ComfyStreamUIPreview:
@@ -9,29 +8,24 @@ class ComfyStreamUIPreview:
     It's needed for ComfyUI to properly register and execute the node.
     The actual implementation is in the JavaScript file.
     """
+
     @classmethod
     def INPUT_TYPES(cls):
-        return {
-            "required": {},
-            "optional": {}
-        }
-
+        return {"required": {}, "optional": {}}
+
     RETURN_TYPES = ()
-    
+
     FUNCTION = "execute"
     CATEGORY = "ComfyStream"
-    
+
     def execute(self):
         # This function doesn't do anything as the real work is done in JavaScript
         # But we need to return something to satisfy the ComfyUI node execution system
         return ("UI Preview Node Executed",)
 
+
 # Register the node class
-NODE_CLASS_MAPPINGS = {
-    "ComfyStreamUIPreview": ComfyStreamUIPreview
-}
+NODE_CLASS_MAPPINGS = {"ComfyStreamUIPreview": ComfyStreamUIPreview}
 
 # Display names for the nodes
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "ComfyStreamUIPreview": "ComfyStream UI Preview"
-} 
+NODE_DISPLAY_NAME_MAPPINGS = {"ComfyStreamUIPreview": "ComfyStream UI Preview"}