Peyton DeNiro
29defc5ef1
Critical fixes: - node_runtime.py: change outer handler in _load_module() from except Exception to except BaseException with sys.modules.pop() cleanup, preventing non-Exception BaseExceptions (SystemExit, KeyboardInterrupt) from violating the no-raise contract - app.py: remove incorrect reassignment of base_node_class_mappings after custom node import; keep it as pre-custom-node baseline so WorkflowPlanner correctly treats extras/custom nodes as NODE_CLASS_MAPPINGS dict lookups instead of direct imports - Regenerate stale test fixtures (text-to-image.py, upscale-model-loader.py) from current source: updated _load_module() with exec failure cleanup pattern, 3-marker structural verification, and try/finally sys.path restoration in import_custom_nodes() Suggestions addressed: - node_runtime.py: _is_comfyui_directory() now checks 3 structural markers (nodes.py, main.py, comfy/) instead of just nodes.py - generator/render.py: added None guard for getattr on generated_helpers.__all__ items with log warning for missing helpers - test_import_path_resolution.py: fixed test_sys_path_restored_on_crash to target the server module specifically after sys.path manipulation, so the try/finally restoration path is actually exercised Nits addressed: - .gitignore: added missing trailing newline New test coverage (tests/test_app_base_mappings.py): - base_node_class_mappings unchanged after custom node import - deep copy independence of base_node_class_mappings - render output contains hardened _load_module with except BaseException pattern
148 lines
4.9 KiB
Python
148 lines
4.9 KiB
Python
import inspect
|
|
import logging
|
|
from pprint import pformat
|
|
from typing import Any
|
|
|
|
import black
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
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, None)
|
|
if func is None:
|
|
log.warning(
|
|
"Helper '%s' missing from generated_helpers — skipping", name
|
|
)
|
|
continue
|
|
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()
|