diff --git a/__init__.py b/__init__.py index 6f3e32fb..b6efc92d 100644 --- a/__init__.py +++ b/__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'] \ No newline at end of file +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/nodes/__init__.py b/nodes/__init__.py index 46aa93c6..0c7e2258 100644 --- a/nodes/__init__.py +++ b/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. +""" diff --git a/nodes/audio_utils/__init__.py b/nodes/audio_utils/__init__.py index 0251eac8..bc090689 100644 --- a/nodes/audio_utils/__init__.py +++ b/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"] diff --git a/nodes/video_stream_utils/__init__.py b/nodes/video_stream_utils/__init__.py index 787f84ce..f4863198 100644 --- a/nodes/video_stream_utils/__init__.py +++ b/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 = {} diff --git a/nodes/web/__init__.py b/nodes/web/__init__.py index ef41b57a..19fc68ea 100644 --- a/nodes/web/__init__.py +++ b/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" -} \ No newline at end of file +NODE_DISPLAY_NAME_MAPPINGS = {"ComfyStreamUIPreview": "ComfyStream UI Preview"} diff --git a/nodes/web/static/.gitkeep b/nodes/web/static/.gitkeep deleted file mode 100644 index e69de29b..00000000