Refactor the exporter into a package-backed module layout.
Split the previous monolithic implementation into focused modules for CLI handling, workflow loading, runtime bootstrap, load ordering, planning, rendering, and file I/O. Keep the public compatibility surface intact by preserving the legacy top-level wrapper while adding as the documented entrypoint. This reduces cross-cutting coupling inside the exporter and makes the runtime/bootstrap behavior easier to evolve without carrying unrelated concerns in one file.
This commit is contained in:
Peyton DeNiro
@@ -0,0 +1,55 @@
|
||||
from typing import TextIO
|
||||
|
||||
from .app import ExportApplication
|
||||
from .cli import DEFAULT_INPUT_FILE, DEFAULT_OUTPUT_FILE, DEFAULT_QUEUE_SIZE, main
|
||||
from .node_runtime import get_node_class_mappings, import_custom_nodes
|
||||
|
||||
|
||||
class ComfyUItoPython:
|
||||
"""Public compatibility facade for the exporter package."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
workflow: str = "",
|
||||
frontend_workflow: str | dict | None = None,
|
||||
input_file: str = "",
|
||||
output_file: str | TextIO = "",
|
||||
queue_size: int = 1,
|
||||
node_class_mappings: dict | None = None,
|
||||
needs_init_custom_nodes: bool = False,
|
||||
):
|
||||
self._app = ExportApplication(
|
||||
workflow=workflow,
|
||||
frontend_workflow=frontend_workflow,
|
||||
input_file=input_file,
|
||||
output_file=output_file,
|
||||
queue_size=queue_size,
|
||||
node_class_mappings=node_class_mappings,
|
||||
needs_init_custom_nodes=needs_init_custom_nodes,
|
||||
node_mapping_loader=get_node_class_mappings,
|
||||
custom_node_importer=import_custom_nodes,
|
||||
)
|
||||
self._app.execute()
|
||||
|
||||
|
||||
def run(
|
||||
input_file: str = DEFAULT_INPUT_FILE,
|
||||
output_file: str = DEFAULT_OUTPUT_FILE,
|
||||
queue_size: int = DEFAULT_QUEUE_SIZE,
|
||||
) -> None:
|
||||
"""Generate Python code from a ComfyUI workflow_api.json file."""
|
||||
ComfyUItoPython(
|
||||
input_file=input_file,
|
||||
output_file=output_file,
|
||||
queue_size=queue_size,
|
||||
needs_init_custom_nodes=True,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ComfyUItoPython",
|
||||
"run",
|
||||
"main",
|
||||
"get_node_class_mappings",
|
||||
"import_custom_nodes",
|
||||
]
|
||||
@@ -0,0 +1,5 @@
|
||||
from .cli import main
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,75 @@
|
||||
import copy
|
||||
|
||||
from typing import TextIO
|
||||
|
||||
from .generator.planner import WorkflowPlanner
|
||||
from .generator.render import WorkflowRenderer
|
||||
from .io import write_python_output
|
||||
from .load_order import LoadOrderDeterminer
|
||||
from .workflow_loader import load_frontend_workflow_data, load_workflow_data
|
||||
|
||||
|
||||
class ExportApplication:
|
||||
"""High-level exporter orchestration."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
workflow: str = "",
|
||||
frontend_workflow: str | dict | None = None,
|
||||
input_file: str = "",
|
||||
output_file: str | TextIO = "",
|
||||
queue_size: int = 1,
|
||||
node_class_mappings: dict | None = None,
|
||||
needs_init_custom_nodes: bool = False,
|
||||
node_mapping_loader=None,
|
||||
custom_node_importer=None,
|
||||
):
|
||||
if input_file and workflow:
|
||||
raise ValueError("Can't provide both input_file and workflow")
|
||||
if not input_file and not workflow:
|
||||
raise ValueError("Needs input_file or workflow")
|
||||
if not output_file:
|
||||
raise ValueError("Needs output_file")
|
||||
|
||||
self.workflow = workflow
|
||||
self.frontend_workflow = frontend_workflow
|
||||
self.input_file = input_file
|
||||
self.output_file = output_file
|
||||
self.queue_size = queue_size
|
||||
self.node_mapping_loader = node_mapping_loader
|
||||
self.custom_node_importer = custom_node_importer
|
||||
self.node_class_mappings = (
|
||||
node_class_mappings
|
||||
if node_class_mappings is not None
|
||||
else self.node_mapping_loader()
|
||||
)
|
||||
self.needs_init_custom_nodes = needs_init_custom_nodes
|
||||
self.base_node_class_mappings = copy.deepcopy(self.node_class_mappings)
|
||||
|
||||
def execute(self) -> None:
|
||||
data = load_workflow_data(self.workflow, self.input_file)
|
||||
metadata_workflow_data = load_frontend_workflow_data(self.frontend_workflow)
|
||||
|
||||
missing_node_types = {
|
||||
node_data["class_type"]
|
||||
for node_data in data.values()
|
||||
if node_data["class_type"] not in self.node_class_mappings
|
||||
}
|
||||
if self.needs_init_custom_nodes or missing_node_types:
|
||||
self.custom_node_importer()
|
||||
self.base_node_class_mappings = copy.deepcopy(self.node_class_mappings)
|
||||
|
||||
load_order = LoadOrderDeterminer(
|
||||
data, self.node_class_mappings
|
||||
).determine_load_order()
|
||||
plan = WorkflowPlanner(
|
||||
self.node_class_mappings, self.base_node_class_mappings
|
||||
).build_plan(
|
||||
load_order,
|
||||
data,
|
||||
metadata_workflow_data,
|
||||
queue_size=self.queue_size,
|
||||
)
|
||||
generated_code = WorkflowRenderer().render(plan)
|
||||
write_python_output(self.output_file, generated_code)
|
||||
print(f"Code successfully generated and written to {self.output_file}")
|
||||
@@ -0,0 +1,42 @@
|
||||
from argparse import ArgumentParser
|
||||
|
||||
DEFAULT_INPUT_FILE = "workflow_api.json"
|
||||
DEFAULT_OUTPUT_FILE = "workflow_api.py"
|
||||
DEFAULT_QUEUE_SIZE = 10
|
||||
|
||||
|
||||
def build_argument_parser() -> ArgumentParser:
|
||||
parser = ArgumentParser(
|
||||
description="Generate Python code from a ComfyUI workflow_api.json file."
|
||||
)
|
||||
parser.add_argument(
|
||||
"-f",
|
||||
"--input_file",
|
||||
type=str,
|
||||
help="path to the input JSON file",
|
||||
default=DEFAULT_INPUT_FILE,
|
||||
)
|
||||
parser.add_argument(
|
||||
"-o",
|
||||
"--output_file",
|
||||
type=str,
|
||||
help="path to the output Python file",
|
||||
default=DEFAULT_OUTPUT_FILE,
|
||||
)
|
||||
parser.add_argument(
|
||||
"-q",
|
||||
"--queue_size",
|
||||
type=int,
|
||||
help="number of times the workflow will be executed by default",
|
||||
default=DEFAULT_QUEUE_SIZE,
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def main() -> None:
|
||||
from . import run
|
||||
|
||||
parser = build_argument_parser()
|
||||
pargs = parser.parse_args()
|
||||
run(**vars(pargs))
|
||||
print("Done.")
|
||||
@@ -0,0 +1,4 @@
|
||||
from .planner import WorkflowPlanner
|
||||
from .render import WorkflowRenderer
|
||||
|
||||
__all__ = ["WorkflowPlanner", "WorkflowRenderer"]
|
||||
@@ -0,0 +1,17 @@
|
||||
from ..node_runtime import (
|
||||
add_comfyui_directory_to_sys_path,
|
||||
add_extra_model_paths,
|
||||
bootstrap_comfyui_runtime,
|
||||
find_path,
|
||||
get_comfyui_path,
|
||||
get_value_at_index,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"add_comfyui_directory_to_sys_path",
|
||||
"add_extra_model_paths",
|
||||
"bootstrap_comfyui_runtime",
|
||||
"find_path",
|
||||
"get_comfyui_path",
|
||||
"get_value_at_index",
|
||||
]
|
||||
@@ -0,0 +1,12 @@
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GenerationPlan:
|
||||
import_statements: dict[str, set[str]]
|
||||
special_functions_code: list[str]
|
||||
loop_code: list[str]
|
||||
workflow_data: dict
|
||||
metadata_workflow_data: dict | None
|
||||
queue_size: int
|
||||
custom_nodes: bool
|
||||
@@ -0,0 +1,228 @@
|
||||
import inspect
|
||||
import json
|
||||
import keyword
|
||||
import random
|
||||
import re
|
||||
from typing import Any, Callable
|
||||
|
||||
from .model import GenerationPlan
|
||||
|
||||
|
||||
class WorkflowPlanner:
|
||||
"""Convert ordered workflow nodes into a renderer-ready generation plan."""
|
||||
|
||||
def __init__(self, node_class_mappings: dict, base_node_class_mappings: dict):
|
||||
self.node_class_mappings = node_class_mappings
|
||||
self.base_node_class_mappings = base_node_class_mappings
|
||||
|
||||
@staticmethod
|
||||
def sanitize_node_id(node_id: str) -> str:
|
||||
sanitized = re.sub(r"[^a-z0-9_]", "_", str(node_id).lower().strip())
|
||||
sanitized = re.sub(r"_+", "_", sanitized).strip("_")
|
||||
if not sanitized:
|
||||
sanitized = "node"
|
||||
return sanitized
|
||||
|
||||
@staticmethod
|
||||
def clean_variable_name(class_type: str) -> str:
|
||||
clean_name = class_type.lower().strip().replace("-", "_").replace(" ", "_")
|
||||
clean_name = re.sub(r"[^a-z0-9_]", "", clean_name)
|
||||
if clean_name[0].isdigit():
|
||||
clean_name = "_" + clean_name
|
||||
return clean_name
|
||||
|
||||
def build_plan(
|
||||
self,
|
||||
load_order: list[tuple[str, dict, bool]],
|
||||
workflow_data: dict,
|
||||
metadata_workflow_data: dict | None = None,
|
||||
queue_size: int = 10,
|
||||
) -> GenerationPlan:
|
||||
import_statements = {"nodes": {"NODE_CLASS_MAPPINGS"}}
|
||||
executed_variables = {}
|
||||
special_functions_code = []
|
||||
code = []
|
||||
initialized_objects = {}
|
||||
custom_nodes = False
|
||||
|
||||
for idx, data, is_special_function in load_order:
|
||||
inputs, class_type = data["inputs"], data["class_type"]
|
||||
input_types = self.node_class_mappings[class_type].INPUT_TYPES()
|
||||
class_def = self.node_class_mappings[class_type]()
|
||||
|
||||
missing_required_variable = False
|
||||
if "required" in input_types.keys():
|
||||
for required in input_types["required"]:
|
||||
if required not in inputs.keys():
|
||||
missing_required_variable = True
|
||||
if missing_required_variable:
|
||||
continue
|
||||
|
||||
if class_type not in initialized_objects:
|
||||
if class_type == "PreviewImage":
|
||||
continue
|
||||
|
||||
class_type, import_statement, class_code = self.get_class_info(
|
||||
class_type
|
||||
)
|
||||
initialized_objects[class_type] = self.clean_variable_name(class_type)
|
||||
if class_type 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)
|
||||
|
||||
class_def_params = self.get_function_parameters(
|
||||
getattr(class_def, class_def.FUNCTION)
|
||||
)
|
||||
no_params = class_def_params is None
|
||||
|
||||
inputs = {
|
||||
key: value
|
||||
for key, value in inputs.items()
|
||||
if no_params or key in class_def_params
|
||||
}
|
||||
|
||||
hidden_inputs = input_types.get("hidden", {})
|
||||
if (
|
||||
"unique_id" in hidden_inputs
|
||||
and (no_params or "unique_id" in class_def_params)
|
||||
):
|
||||
inputs["unique_id"] = random.randint(1, 2**64)
|
||||
if "prompt" in hidden_inputs and (no_params or "prompt" in class_def_params):
|
||||
inputs["prompt"] = {"variable_name": "prompt"}
|
||||
if "extra_pnginfo" in hidden_inputs and (
|
||||
no_params or "extra_pnginfo" in class_def_params
|
||||
):
|
||||
inputs["extra_pnginfo"] = {"variable_name": "extra_pnginfo"}
|
||||
if "hidden" not in input_types and class_def_params is not None:
|
||||
if "unique_id" in class_def_params:
|
||||
inputs["unique_id"] = random.randint(1, 2**64)
|
||||
|
||||
executed_variables[idx] = (
|
||||
f"{self.clean_variable_name(class_type)}_"
|
||||
f"{self.sanitize_node_id(str(idx))}"
|
||||
)
|
||||
inputs = self.update_inputs(inputs, executed_variables)
|
||||
seed_sync_code = self.create_prompt_seed_sync_code(
|
||||
idx, inputs, is_special_function
|
||||
)
|
||||
|
||||
target_lines = special_functions_code if is_special_function else code
|
||||
if seed_sync_code:
|
||||
target_lines.extend(seed_sync_code)
|
||||
target_lines.append(
|
||||
self.create_function_call_code(
|
||||
initialized_objects[class_type],
|
||||
class_def.FUNCTION,
|
||||
executed_variables[idx],
|
||||
is_special_function,
|
||||
**inputs,
|
||||
)
|
||||
)
|
||||
|
||||
return GenerationPlan(
|
||||
import_statements=import_statements,
|
||||
special_functions_code=special_functions_code,
|
||||
loop_code=code,
|
||||
workflow_data=workflow_data,
|
||||
metadata_workflow_data=metadata_workflow_data,
|
||||
queue_size=queue_size,
|
||||
custom_nodes=custom_nodes,
|
||||
)
|
||||
|
||||
def create_function_call_code(
|
||||
self,
|
||||
obj_name: str,
|
||||
func: str,
|
||||
variable_name: str,
|
||||
is_special_function: bool,
|
||||
**kwargs,
|
||||
) -> str:
|
||||
args = ", ".join(self.format_arg(key, value) for key, value in kwargs.items())
|
||||
code = f"{variable_name} = {obj_name}.{func}({args})\n"
|
||||
if not is_special_function:
|
||||
code = f"\t{code}"
|
||||
return code
|
||||
|
||||
def create_prompt_seed_sync_code(
|
||||
self, node_id: str, inputs: dict, is_special_function: bool
|
||||
) -> list[str]:
|
||||
seed_sync_lines = []
|
||||
for key in ("seed", "noise_seed"):
|
||||
if key not in inputs:
|
||||
continue
|
||||
randomized_seed_variable = (
|
||||
f"node_{self.sanitize_node_id(str(node_id))}_{self.clean_variable_name(key)}"
|
||||
)
|
||||
seed_sync_lines.append(
|
||||
f'{randomized_seed_variable} = prompt["{node_id}"]["inputs"]["{key}"] = random.randint(1, 2**64)'
|
||||
)
|
||||
inputs[key] = {"variable_name": randomized_seed_variable}
|
||||
|
||||
if not seed_sync_lines:
|
||||
return []
|
||||
|
||||
indentation = "" if is_special_function else "\t"
|
||||
return [f"{indentation}{line}\n" for line in seed_sync_lines]
|
||||
|
||||
def format_arg(self, key: str, value: Any) -> str:
|
||||
value_code = self.format_arg_value(key, value)
|
||||
if key.isidentifier() and not keyword.iskeyword(key):
|
||||
return f"{key}={value_code}"
|
||||
return f"**{{{json.dumps(key)}: {value_code}}}"
|
||||
|
||||
@staticmethod
|
||||
def format_arg_value(key: str, value: Any) -> str:
|
||||
if isinstance(value, dict) and "variable_name" in value:
|
||||
return value["variable_name"]
|
||||
if key == "noise_seed" or key == "seed":
|
||||
return "random.randint(1, 2**64)"
|
||||
if isinstance(value, str):
|
||||
return json.dumps(value)
|
||||
return repr(value)
|
||||
|
||||
def get_class_info(self, class_type: str) -> tuple[str, tuple[str, str], str]:
|
||||
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)
|
||||
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
|
||||
|
||||
@staticmethod
|
||||
def get_function_parameters(func: Callable) -> list | None:
|
||||
signature = inspect.signature(func)
|
||||
parameters = {
|
||||
name: param.default if param.default != param.empty else None
|
||||
for name, param in signature.parameters.items()
|
||||
}
|
||||
catch_all = any(
|
||||
param.kind == inspect.Parameter.VAR_KEYWORD
|
||||
for param in signature.parameters.values()
|
||||
)
|
||||
return list(parameters.keys()) if not catch_all else None
|
||||
|
||||
def update_inputs(self, inputs: dict, executed_variables: dict) -> dict:
|
||||
for key in inputs.keys():
|
||||
if (
|
||||
isinstance(inputs[key], list)
|
||||
and inputs[key][0] in executed_variables.keys()
|
||||
):
|
||||
inputs[key] = {
|
||||
"variable_name": f"get_value_at_index({executed_variables[inputs[key][0]]}, {inputs[key][1]})"
|
||||
}
|
||||
return inputs
|
||||
@@ -0,0 +1,135 @@
|
||||
import inspect
|
||||
from pprint import pformat
|
||||
from typing import Any
|
||||
|
||||
import black
|
||||
|
||||
from ..node_runtime import import_custom_nodes
|
||||
from .generated_helpers import (
|
||||
add_comfyui_directory_to_sys_path,
|
||||
add_extra_model_paths,
|
||||
bootstrap_comfyui_runtime,
|
||||
find_path,
|
||||
get_comfyui_path,
|
||||
get_value_at_index,
|
||||
)
|
||||
from .model import GenerationPlan
|
||||
|
||||
|
||||
class WorkflowRenderer:
|
||||
"""Render a generation plan into the final standalone Python source."""
|
||||
|
||||
def render(self, plan: GenerationPlan) -> str:
|
||||
workflow_literal = self.format_python_literal(plan.workflow_data)
|
||||
if plan.metadata_workflow_data is None:
|
||||
extra_pnginfo_literal = "None"
|
||||
else:
|
||||
extra_pnginfo_literal = self.format_python_literal(
|
||||
{"workflow": plan.metadata_workflow_data}
|
||||
)
|
||||
|
||||
func_strings = []
|
||||
for func in [
|
||||
get_value_at_index,
|
||||
get_comfyui_path,
|
||||
find_path,
|
||||
add_comfyui_directory_to_sys_path,
|
||||
add_extra_model_paths,
|
||||
bootstrap_comfyui_runtime,
|
||||
]:
|
||||
func_strings.append(f"\n{inspect.getsource(func)}")
|
||||
|
||||
static_imports = [
|
||||
"# Imports",
|
||||
"import json",
|
||||
"import os",
|
||||
"import random",
|
||||
"import sys",
|
||||
"from typing import Sequence, Mapping, Any, Union",
|
||||
] + func_strings
|
||||
static_imports.append("\n# Runtime support")
|
||||
static_imports.append(
|
||||
"\nbootstrap_comfyui_runtime()\nimport torch\nadd_extra_model_paths()\n"
|
||||
)
|
||||
|
||||
if plan.custom_nodes:
|
||||
static_imports.append(f"\n{inspect.getsource(import_custom_nodes)}\n")
|
||||
custom_nodes_call = "import_custom_nodes()"
|
||||
else:
|
||||
custom_nodes_call = None
|
||||
|
||||
imports_code = []
|
||||
for module_name in sorted(plan.import_statements.keys()):
|
||||
class_names = ", ".join(sorted(plan.import_statements[module_name]))
|
||||
imports_code.append(f"from {module_name} import {class_names}")
|
||||
|
||||
workflow_section = [
|
||||
"# Workflow data",
|
||||
"def build_workflow() -> dict[str, Any]:",
|
||||
f" return {workflow_literal}",
|
||||
"",
|
||||
"def build_extra_pnginfo() -> dict[str, Any] | None:",
|
||||
f" return {extra_pnginfo_literal}",
|
||||
"",
|
||||
"workflow = build_workflow()",
|
||||
"prompt = json.loads(json.dumps(workflow))",
|
||||
"extra_pnginfo = build_extra_pnginfo()",
|
||||
]
|
||||
|
||||
execution_section = ["# Workflow execution"]
|
||||
if imports_code:
|
||||
execution_section.extend(["# Node imports", *imports_code, ""])
|
||||
execution_section.extend(
|
||||
[
|
||||
"def main():",
|
||||
" with torch.inference_mode():",
|
||||
]
|
||||
)
|
||||
if custom_nodes_call:
|
||||
execution_section.append(f" {custom_nodes_call}")
|
||||
execution_section.extend(
|
||||
self.build_function_body(
|
||||
plan.special_functions_code, "pass", indentation=" "
|
||||
).splitlines()
|
||||
)
|
||||
execution_section.append(f" for q in range({plan.queue_size}):")
|
||||
execution_section.extend(
|
||||
self.build_function_body(
|
||||
plan.loop_code, "pass", indentation=" "
|
||||
).splitlines()
|
||||
)
|
||||
|
||||
entrypoint_section = [
|
||||
"# Entrypoint",
|
||||
'if __name__ == "__main__":',
|
||||
" main()",
|
||||
]
|
||||
|
||||
final_code = "\n".join(
|
||||
static_imports
|
||||
+ [""]
|
||||
+ workflow_section
|
||||
+ [""]
|
||||
+ execution_section
|
||||
+ [""]
|
||||
+ entrypoint_section
|
||||
)
|
||||
return black.format_str(final_code, mode=black.Mode())
|
||||
|
||||
@staticmethod
|
||||
def format_python_literal(value: Any) -> str:
|
||||
return pformat(value, sort_dicts=False)
|
||||
|
||||
@staticmethod
|
||||
def build_function_body(
|
||||
code_lines: list[str], empty_fallback: str, indentation: str = " "
|
||||
) -> str:
|
||||
if not code_lines:
|
||||
return f"{indentation}{empty_fallback}"
|
||||
formatted_lines = []
|
||||
for line in code_lines:
|
||||
stripped_line = line.lstrip()
|
||||
if not stripped_line.endswith("\n"):
|
||||
stripped_line += "\n"
|
||||
formatted_lines.append(f"{indentation}{stripped_line}")
|
||||
return "".join(formatted_lines).rstrip()
|
||||
@@ -0,0 +1,23 @@
|
||||
import json
|
||||
import os
|
||||
from typing import TextIO
|
||||
|
||||
|
||||
def load_json_input(file_path: str | TextIO, encoding: str = "utf-8") -> dict:
|
||||
"""Read workflow JSON from a file path or file-like object."""
|
||||
if hasattr(file_path, "read"):
|
||||
return json.load(file_path)
|
||||
with open(file_path, "r", encoding=encoding) as file:
|
||||
return json.load(file)
|
||||
|
||||
|
||||
def write_python_output(file_path: str | TextIO, code: str) -> None:
|
||||
"""Write generated Python to a file path or file-like object."""
|
||||
if isinstance(file_path, str):
|
||||
directory = os.path.dirname(file_path)
|
||||
if directory and not os.path.exists(directory):
|
||||
os.makedirs(directory)
|
||||
with open(file_path, "w", encoding="utf-8") as file:
|
||||
file.write(code)
|
||||
return
|
||||
file_path.write(code)
|
||||
@@ -0,0 +1,43 @@
|
||||
from typing import Dict
|
||||
|
||||
|
||||
class LoadOrderDeterminer:
|
||||
"""Determine workflow execution order with loader-like nodes prioritized."""
|
||||
|
||||
def __init__(self, data: Dict, node_class_mappings: Dict):
|
||||
self.data = data
|
||||
self.node_class_mappings = node_class_mappings
|
||||
self.visited = {}
|
||||
self.load_order = []
|
||||
self.is_special_function = False
|
||||
|
||||
def determine_load_order(self) -> list[tuple[str, Dict, bool]]:
|
||||
self._load_special_functions_first()
|
||||
self.is_special_function = False
|
||||
for key in self.data:
|
||||
if key not in self.visited:
|
||||
self._dfs(key)
|
||||
return self.load_order
|
||||
|
||||
def _dfs(self, key: str) -> None:
|
||||
self.visited[key] = True
|
||||
inputs = self.data[key]["inputs"]
|
||||
for value in inputs.values():
|
||||
if isinstance(value, list) and value[0] not in self.visited:
|
||||
self._dfs(value[0])
|
||||
self.load_order.append((key, self.data[key], self.is_special_function))
|
||||
|
||||
def _load_special_functions_first(self) -> None:
|
||||
for key in self.data:
|
||||
class_def = self.node_class_mappings[self.data[key]["class_type"]]()
|
||||
if (
|
||||
class_def.CATEGORY == "loaders"
|
||||
or class_def.FUNCTION in ["encode"]
|
||||
or not any(
|
||||
isinstance(value, list)
|
||||
for value in self.data[key]["inputs"].values()
|
||||
)
|
||||
):
|
||||
self.is_special_function = True
|
||||
if key not in self.visited:
|
||||
self._dfs(key)
|
||||
@@ -0,0 +1,133 @@
|
||||
import os
|
||||
import sys
|
||||
from typing import Any, Mapping, Sequence, Union
|
||||
|
||||
|
||||
def find_path(name: str, path: str = None) -> str:
|
||||
"""Recursively search parent folders until the named entry is found."""
|
||||
if path is None:
|
||||
path = os.getcwd()
|
||||
|
||||
if name in os.listdir(path):
|
||||
path_name = os.path.join(path, name)
|
||||
print(f"{name} found: {path_name}")
|
||||
return path_name
|
||||
|
||||
parent_directory = os.path.dirname(path)
|
||||
if parent_directory == path:
|
||||
return None
|
||||
|
||||
return find_path(name, parent_directory)
|
||||
|
||||
|
||||
def get_comfyui_path() -> str:
|
||||
"""Return the configured ComfyUI path, preferring COMFYUI_PATH when set."""
|
||||
comfyui_path = os.environ.get("COMFYUI_PATH")
|
||||
if comfyui_path:
|
||||
return comfyui_path
|
||||
return find_path("ComfyUI")
|
||||
|
||||
|
||||
def add_comfyui_directory_to_sys_path() -> None:
|
||||
"""Add the ComfyUI checkout to sys.path."""
|
||||
comfyui_path = get_comfyui_path()
|
||||
if comfyui_path is not None and os.path.isdir(comfyui_path):
|
||||
if comfyui_path in sys.path:
|
||||
sys.path.remove(comfyui_path)
|
||||
sys.path.insert(0, comfyui_path)
|
||||
print(f"'{comfyui_path}' added to sys.path")
|
||||
|
||||
|
||||
def add_extra_model_paths() -> None:
|
||||
"""Load ComfyUI extra model paths configuration when available."""
|
||||
try:
|
||||
from main import load_extra_path_config
|
||||
except ImportError:
|
||||
print(
|
||||
"Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead."
|
||||
)
|
||||
from utils.extra_config import load_extra_path_config
|
||||
|
||||
extra_model_paths = find_path("extra_model_paths.yaml")
|
||||
if extra_model_paths is not None:
|
||||
load_extra_path_config(extra_model_paths)
|
||||
else:
|
||||
print("Could not find the extra_model_paths config file.")
|
||||
|
||||
|
||||
def bootstrap_comfyui_runtime() -> None:
|
||||
"""Mirror the allocator-related ComfyUI startup steps before torch import."""
|
||||
add_comfyui_directory_to_sys_path()
|
||||
|
||||
import comfy.options
|
||||
|
||||
comfy.options.enable_args_parsing()
|
||||
|
||||
from comfy.cli_args import args
|
||||
|
||||
if os.name == "nt":
|
||||
os.environ["MIMALLOC_PURGE_DELAY"] = "0"
|
||||
|
||||
if args.default_device is not None:
|
||||
default_dev = args.default_device
|
||||
devices = list(range(32))
|
||||
devices.remove(default_dev)
|
||||
devices.insert(0, default_dev)
|
||||
devices = ",".join(map(str, devices))
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = str(devices)
|
||||
os.environ["HIP_VISIBLE_DEVICES"] = str(devices)
|
||||
|
||||
if args.cuda_device is not None:
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.cuda_device)
|
||||
os.environ["HIP_VISIBLE_DEVICES"] = str(args.cuda_device)
|
||||
os.environ["ASCEND_RT_VISIBLE_DEVICES"] = str(args.cuda_device)
|
||||
|
||||
if args.oneapi_device_selector is not None:
|
||||
os.environ["ONEAPI_DEVICE_SELECTOR"] = args.oneapi_device_selector
|
||||
|
||||
if args.deterministic and "CUBLAS_WORKSPACE_CONFIG" not in os.environ:
|
||||
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
|
||||
|
||||
import cuda_malloc
|
||||
|
||||
if "rocm" in cuda_malloc.get_torch_version_noimport():
|
||||
os.environ["OCL_SET_SVM_SIZE"] = "262144"
|
||||
|
||||
|
||||
def import_custom_nodes() -> None:
|
||||
"""Initialize ComfyUI custom nodes in the exporter runtime."""
|
||||
comfyui_path = get_comfyui_path()
|
||||
if comfyui_path and comfyui_path not in sys.path:
|
||||
sys.path.insert(0, comfyui_path)
|
||||
|
||||
import asyncio
|
||||
import execution
|
||||
from nodes import init_extra_nodes
|
||||
|
||||
if comfyui_path in sys.path:
|
||||
sys.path.remove(comfyui_path)
|
||||
sys.path.insert(0, comfyui_path)
|
||||
|
||||
import server
|
||||
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
server_instance = server.PromptServer(loop)
|
||||
execution.PromptQueue(server_instance)
|
||||
asyncio.run(init_extra_nodes())
|
||||
|
||||
|
||||
def get_node_class_mappings() -> dict:
|
||||
"""Load ComfyUI node mappings on demand."""
|
||||
add_comfyui_directory_to_sys_path()
|
||||
from nodes import NODE_CLASS_MAPPINGS
|
||||
|
||||
return NODE_CLASS_MAPPINGS
|
||||
|
||||
|
||||
def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
|
||||
"""Return a sequence or mapping result item by index."""
|
||||
try:
|
||||
return obj[index]
|
||||
except KeyError:
|
||||
return obj["result"][index]
|
||||
@@ -0,0 +1,19 @@
|
||||
import json
|
||||
|
||||
from .io import load_json_input
|
||||
|
||||
|
||||
def load_workflow_data(workflow: str, input_file: str):
|
||||
"""Load workflow data from inline JSON or an input file."""
|
||||
if input_file:
|
||||
return load_json_input(input_file)
|
||||
return json.loads(workflow)
|
||||
|
||||
|
||||
def load_frontend_workflow_data(frontend_workflow: str | dict | None):
|
||||
"""Load optional frontend workflow metadata."""
|
||||
if not frontend_workflow:
|
||||
return None
|
||||
if isinstance(frontend_workflow, str):
|
||||
return json.loads(frontend_workflow)
|
||||
return frontend_workflow
|
||||
Reference in New Issue
Block a user