Hotfix for node class names with trailing white spaces. Removed the need to import functions from the utils file when running the generated code.

Author: Peyton

comfyui_to_python.py

@@ -10,7 +10,7 @@
 
 from typing import Dict, List, Any, Callable, Tuple
 
-from utils import import_custom_nodes, add_comfyui_directories_to_sys_path
+from utils import import_custom_nodes, add_comfyui_directory_to_sys_path, get_value_at_index
 
 
 sys.path.append('../')
@@ -218,7 +218,7 @@ def generate_workflow(self, load_order: List, filename: str = 'generated_code_wo
                     continue
 
                 class_type, import_statement, class_code = self.get_class_info(class_type)
-                initialized_objects[class_type] = class_type.lower()
+                initialized_objects[class_type] = class_type.lower().strip()
                 if class_type in self.base_node_class_mappings.keys():
                     import_statements.add(import_statement)
                 if class_type not in self.base_node_class_mappings.keys():
@@ -232,7 +232,7 @@ def generate_workflow(self, load_order: List, filename: str = 'generated_code_wo
             inputs = {key: value for key, value in inputs.items() if key in class_def_params}
 
             # Create executed variable and generate code
-            executed_variables[idx] = f'{class_type.lower()}_{idx}'
+            executed_variables[idx] = f'{class_type.lower().strip()}_{idx}'
             inputs = self.update_inputs(inputs, executed_variables)
 
             if is_special_function:
@@ -302,14 +302,16 @@ def assemble_python_code(self, import_statements: set, speical_functions_code: L
         Returns:
             str: Generated final code as a string.
         """
-        # Get the source code of the function as a string
-        add_comfyui_directories_to_sys_path_code = inspect.getsource(add_comfyui_directories_to_sys_path)
+        # Get the source code of the utils functions as a string
+        func_strings = []
+        for func in [add_comfyui_directory_to_sys_path, get_value_at_index]:
+            func_strings.append(f'\n{inspect.getsource(func)}')
         # Define static import statements required for the script
-        static_imports = ['import os', 'import random', 'import sys', 'import torch', f'\n{add_comfyui_directories_to_sys_path_code}', 
-                        '\n\nadd_comfyui_directories_to_sys_path()']
+        static_imports = ['import os', 'import random', 'import sys', 'from typing import Sequence, Mapping, Any, Union', 
+                          'import torch'] + func_strings + ['\n\nadd_comfyui_directory_to_sys_path()']
         # Check if custom nodes should be included
         if custom_nodes:
-            static_imports.append('\nfrom utils import import_custom_nodes, get_value_at_index\n')
+            static_imports.append(f'\n{inspect.getsource(import_custom_nodes)}\n')
             custom_nodes = 'import_custom_nodes()\n\t'
         else:
             custom_nodes = ''
@@ -336,9 +338,9 @@ def get_class_info(self, class_type: str) -> Tuple[str, str, str]:
         """
         import_statement = class_type
         if class_type in self.base_node_class_mappings.keys():
-            class_code = f'{class_type.lower()} = {class_type}()'
+            class_code = f'{class_type.lower().strip()} = {class_type.strip()}()'
         else:
-            class_code = f'{class_type.lower()} = NODE_CLASS_MAPPINGS["{class_type}"]()'
+            class_code = f'{class_type.lower().strip()} = NODE_CLASS_MAPPINGS["{class_type}"]()'
 
         return class_type, import_statement, class_code
 

utils.py

@@ -1,22 +1,19 @@
-import asyncio
-import json
-import glob
 import os
 from typing import Sequence, Mapping, Any, Union
 import sys
 
 sys.path.append('../')
 
-import execution
-from nodes import init_custom_nodes
-        
 
 def import_custom_nodes() -> None:
     """Find all custom nodes in the custom_nodes folder and add those node objects to NODE_CLASS_MAPPINGS
 
     This function sets up a new asyncio event loop, initializes the PromptServer,
     creates a PromptQueue, and initializes the custom nodes.
     """
+    import asyncio
+    import execution
+    from nodes import init_custom_nodes
     import server
 
     # Creating a new event loop and setting it as the default loop
@@ -30,21 +27,19 @@ def import_custom_nodes() -> None:
     # Initializing custom nodes
     init_custom_nodes()
 
-
-def add_comfyui_directories_to_sys_path() -> None:
+def add_comfyui_directory_to_sys_path() -> None:
     """
-    Recursively looks at parent folders starting from the current working directory until it finds 'ComfyUI' and 'ComfyUI-to-Python-Extension'.
-    Once found, the directories are added to sys.path.
+    Recursively looks at parent folders starting from the current working directory until it finds 'ComfyUI'.
+    Once found, the directory is added to sys.path.
     """
     start_path = os.getcwd()  # Get the current working directory
 
     def search_directory(path: str) -> None:
-        # Check if the current directory contains 'ComfyUI' or 'ComfyUI-to-Python-Extension'
-        for directory_name in ['ComfyUI', 'ComfyUI-to-Python-Extension']:
-            if directory_name in os.listdir(path):
-                directory_path = os.path.join(path, directory_name)
-                sys.path.append(directory_path)
-                print(f"'{directory_name}' found and added to sys.path: {directory_path}")
+        # Check if the current directory contains 'ComfyUI'
+        if 'ComfyUI' in os.listdir(path):
+            directory_path = os.path.join(path, 'ComfyUI')
+            sys.path.append(directory_path)
+            print(f"ComfyUI found and added to sys.path: {directory_path}")
 
         # Get the parent directory
         parent_directory = os.path.dirname(path)