nodes/pydn-ComfyUI-to-Python-Extension
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix upscale export runtime compatibility

Browse Source
This commit is contained in:
Peyton DeNiro
2026-03-28 13:08:34 -05:00
parent 72a87c31b3
commit 557adae33e
1 changed files with 52 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
comfyui_to_python.py
+52 -19
View File
@@ -16,20 +16,30 @@ sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from comfyui_to_python_utils import (
import_custom_nodes,
find_path,
get_comfyui_path,
add_comfyui_directory_to_sys_path,
add_extra_model_paths,
get_value_at_index,
)
add_comfyui_directory_to_sys_path()
from nodes import NODE_CLASS_MAPPINGS
DEFAULT_INPUT_FILE = "workflow_api.json"
DEFAULT_OUTPUT_FILE = "workflow_api.py"
DEFAULT_QUEUE_SIZE = 10
def get_node_class_mappings() -> Dict:
"""Load ComfyUI node mappings on demand.
Tests that inject explicit node mappings should not need a full ComfyUI runtime
just to import this module.
"""
add_comfyui_directory_to_sys_path()
from nodes import NODE_CLASS_MAPPINGS
return NODE_CLASS_MAPPINGS
class FileHandler:
"""Handles reading and writing files.
@@ -199,7 +209,7 @@ class CodeGenerator:
"""
# Create the necessary data structures to hold imports and generated code
import_statements, executed_variables, special_functions_code, code = (
set(["NODE_CLASS_MAPPINGS"]),
{"nodes": {"NODE_CLASS_MAPPINGS"}},
{},
[],
[],
@@ -235,8 +245,9 @@ class CodeGenerator:
)
initialized_objects[class_type] = self.clean_variable_name(class_type)
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():
module_name, import_name = import_statement
import_statements.setdefault(module_name, set()).add(import_name)
if 'NODE_CLASS_MAPPINGS["' in class_code:
custom_nodes = True
special_functions_code.append(class_code)
@@ -289,7 +300,11 @@ class CodeGenerator:
# Generate final code by combining imports and code, and wrap them in a main function
final_code = self.assemble_python_code(
import_statements, special_functions_code, code, queue_size, custom_nodes
import_statements,
special_functions_code,
code,
queue_size,
custom_nodes,
)
return final_code
@@ -347,7 +362,7 @@ class CodeGenerator:
def assemble_python_code(
self,
import_statements: set,
import_statements: Dict[str, set],
speical_functions_code: List[str],
code: List[str],
queue_size: int,
@@ -356,7 +371,7 @@ class CodeGenerator:
"""Generates the final code string.
Args:
import_statements (set): A set of unique import statements.
import_statements (Dict[str, set]): Import statements grouped by module.
speical_functions_code (List[str]): A list of special functions code strings.
code (List[str]): A list of code strings.
queue_size (int): Number of photos that will be generated by the script.
@@ -369,6 +384,7 @@ class CodeGenerator:
func_strings = []
for func in [
get_value_at_index,
get_comfyui_path,
find_path,
add_comfyui_directory_to_sys_path,
add_extra_model_paths,
@@ -393,9 +409,10 @@ class CodeGenerator:
else:
custom_nodes = ""
# Create import statements for node classes
imports_code = [
f"from nodes import {', '.join([class_name for class_name in import_statements])}"
]
imports_code = []
for module_name in sorted(import_statements.keys()):
class_names = ", ".join(sorted(import_statements[module_name]))
imports_code.append(f"from {module_name} import {class_names}")
# Assemble the main function code, including custom nodes if applicable
main_function_code = (
"def main():\n\t"
@@ -415,20 +432,31 @@ class CodeGenerator:
return final_code
def get_class_info(self, class_type: str) -> Tuple[str, str, str]:
def get_class_info(self, class_type: str) -> Tuple[str, Tuple[str, str], str]:
"""Generates and returns necessary information about class type.
Args:
class_type (str): Class type.
Returns:
Tuple[str, str, str]: Updated class type, import statement string, class initialization code.
Tuple[str, Tuple[str, str], str]: Updated class type, import statement and initialization code.
"""
import_statement = class_type
class_obj = self.base_node_class_mappings.get(class_type)
module_name = "nodes"
if class_obj is not None:
module_name = class_obj.__module__
variable_name = self.clean_variable_name(class_type)
if class_type in self.base_node_class_mappings.keys():
is_importable_module = bool(
module_name
and "/" not in module_name
and "\\" not in module_name
and all(part.isidentifier() for part in module_name.split("."))
)
if class_type in self.base_node_class_mappings.keys() and is_importable_module:
import_statement = (module_name, class_type)
class_code = f"{variable_name} = {class_type.strip()}()"
else:
import_statement = ("nodes", "NODE_CLASS_MAPPINGS")
class_code = f'{variable_name} = NODE_CLASS_MAPPINGS["{class_type}"]()'
return class_type, import_statement, class_code
@@ -514,7 +542,7 @@ class ComfyUItoPython:
input_file: str = "",
output_file: str | TextIO = "",
queue_size: int = 1,
node_class_mappings: Dict = NODE_CLASS_MAPPINGS,
node_class_mappings: Dict | None = None,
needs_init_custom_nodes: bool = False,
):
"""Initialize the ComfyUItoPython class with the given parameters. Exactly one of workflow or input_file must be specified.
@@ -523,7 +551,8 @@ class ComfyUItoPython:
input_file (str): Path to the input JSON file.
output_file (str | TextIO): Path to the output file or a file-like object.
queue_size (int): The number of times a workflow will be executed by the script. Defaults to 1.
node_class_mappings (Dict): Mappings of node classes. Defaults to NODE_CLASS_MAPPINGS.
node_class_mappings (Dict | None): Mappings of node classes. Defaults to the current
ComfyUI NODE_CLASS_MAPPINGS when not provided.
needs_init_custom_nodes (bool): Whether to initialize custom nodes. Defaults to False.
"""
if input_file and workflow:
@@ -538,7 +567,11 @@ class ComfyUItoPython:
self.input_file = input_file
self.output_file = output_file
self.queue_size = queue_size
self.node_class_mappings = node_class_mappings
self.node_class_mappings = (
node_class_mappings
if node_class_mappings is not None
else get_node_class_mappings()
)
self.needs_init_custom_nodes = needs_init_custom_nodes
self.base_node_class_mappings = copy.deepcopy(self.node_class_mappings)