Hotfix for node class names with trailing white spaces. Removed the n…

…eed to import functions from the utils file when running the generated code.

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)