Peyton DeNiro
ae8d5a2f2a
Critical:
- _load_module(): wrap exec_module() in try/except, pop from sys.modules
on failure so broken modules aren't cached
- get_comfyui_path(): add _is_comfyui_directory() check to third (CWD walk)
strategy so a non-ComfyUI directory named "ComfyUI" isn't accepted
- get_node_class_mappings(): return cached sys.modules["nodes"] if already
loaded to avoid re-executing nodes.py and losing init_extra_nodes() state
Suggestions:
- import_custom_nodes(): wrap sys.path filtering in try/finally so path is
always restored even if _load_module("server") raises mid-execution
- render.py: auto-discover helpers from generated_helpers.__all__ instead of
a manually maintained list
- pyproject.toml: remove duplicate ruff dev dependency (kept standard
[project.optional-dependencies] form)
Nitpicks:
- Remove dead imports (types, StringIO) from test_import_path_resolution.py
- Trim trailing blank lines in pyproject.toml
- Tone down _load_module docstring: "eliminates ALL" -> "significantly reduces"
Tests (7 new cases):
- TestLoadModuleFailureCleanup: broken module removed on exec failure, retry
starts fresh
- TestLoadModuleTemp: temp modules purged from sys.modules (success + failure)
- TestRepeatedNodeLoad: repeated _load_module("nodes") returns cached copy
- TestGetComfyuiPathThirdStrategy: fake "ComfyUI" dir rejected
- TestImportCustomNodesSysPathRestoration: sys.path restored after crash
140 lines
4.6 KiB
Python
140 lines
4.6 KiB
Python
import inspect
|
|
from pprint import pformat
|
|
from typing import Any
|
|
|
|
import black
|
|
|
|
from ..node_runtime import import_custom_nodes
|
|
from . import generated_helpers
|
|
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}
|
|
)
|
|
|
|
# Auto-discover helpers from generated_helpers.__all__ so new helpers
|
|
# are picked up without manually updating this list.
|
|
func_strings = []
|
|
for name in generated_helpers.__all__:
|
|
func = getattr(generated_helpers, name)
|
|
func_strings.append(f"\n{inspect.getsource(func)}")
|
|
|
|
static_imports = [
|
|
"# Imports",
|
|
"import importlib.util",
|
|
"import json",
|
|
"import logging",
|
|
"import os",
|
|
"import random",
|
|
"import sys",
|
|
"from typing import Sequence, Mapping, Any, Union",
|
|
"",
|
|
"log = logging.getLogger(__name__)",
|
|
] + func_strings
|
|
|
|
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",
|
|
"def main(unload_models: bool | None = None):",
|
|
" bootstrap_comfyui_runtime()",
|
|
" add_extra_model_paths()",
|
|
]
|
|
if custom_nodes_call:
|
|
execution_section.append(f" {custom_nodes_call}")
|
|
if imports_code:
|
|
execution_section.extend(["", " # Node imports"])
|
|
execution_section.extend(f" {line}" for line in imports_code)
|
|
execution_section.extend(
|
|
[
|
|
"",
|
|
" import torch",
|
|
"",
|
|
" try:",
|
|
" with torch.inference_mode():",
|
|
]
|
|
)
|
|
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()
|
|
)
|
|
execution_section.extend(
|
|
[
|
|
" finally:",
|
|
" cleanup_comfyui_runtime(unload_models=unload_models)",
|
|
]
|
|
)
|
|
|
|
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()
|