Compare commits

...
9 Commits
Author SHA1 Message Date
Peyton DeNiro 6cdcc235a0 update README to add additional troubleshooting steps for installation 2026-04-11 19:53:04 -05:00
Peyton DeNiro 9a17000417 Update version. 2026-03-29 19:06:52 -05:00
Peyton DeNiro fbc06463f7 Merge pull request #153 from pydn/runtime-parity-cache-memory-lifecycle
Add generated runtime cleanup after workflow execution
2026-03-29 19:02:51 -05:00
Peyton DeNiro 77c0bbec9b Guard ComfyUI cleanup teardown failures 2026-03-29 19:02:21 -05:00
Peyton DeNiro f9163f2375 Add generated runtime cleanup after workflow execution. Exported scripts now run best-effort ComfyUI cache cleanup in a finally block and support opt-in model unloading for embedded repeated-call use cases. Document the supported VRAM-related runtime flags and lock the behavior in focused exporter and runtime cleanup tests. 2026-03-29 18:18:37 -05:00
Peyton DeNiro d6c5b110ab Merge pull request #152 from pydn/cluster-2-compatibility-78-99-137
Clarify Python support and harden export/install compatibility
2026-03-29 17:55:19 -05:00
Peyton DeNiro 5d70efdc88 Align Cluster 2 compatibility boundaries
Declare Python 3.12 support, replace the hidden pip bootstrap with explicit uv sync guidance, and remove prompt() from the Web UI save flow. Add focused contract coverage for the metadata and frontend defaults.
2026-03-29 17:34:30 -05:00
Peyton DeNiro 228c3a076a Merge pull request #151 from pydn/generator-codegen-correctness
Make generator codegen regressions explicit and fix string seed rand
2026-03-29 17:13:10 -05:00
Peyton DeNiro 574d96369c Make generator codegen regressions explicit and fix string seed randomization 2026-03-29 17:09:30 -05:00
20 changed files with 874 additions and 116 deletions
+53 -10
View File
@@ -13,22 +13,20 @@ This project supports:
## Install
Choose the setup that matches how you want to use the project.
This project supports Python 3.12 and newer.
### Web UI extension (`File -> Save As Script`)
For ComfyUI to recognize this project as an extension, the repo must be discoverable through ComfyUI's `custom_nodes` search paths.
Use one of these setups:
Clone directly into `ComfyUI/custom_nodes`:
1. Clone directly into `ComfyUI/custom_nodes`
```bash
cd /path/to/ComfyUI/custom_nodes
git clone https://github.com/pydn/ComfyUI-to-Python-Extension.git
cd ComfyUI-to-Python-Extension
uv sync
```
2. Keep the repo elsewhere, then either:
Or keep the repo elsewhere, then either:
- symlink it into `ComfyUI/custom_nodes`
- add its parent directory to ComfyUI's `custom_nodes` search paths via `extra_model_paths.yaml`
@@ -37,10 +35,29 @@ Example symlink setup:
git clone https://github.com/pydn/ComfyUI-to-Python-Extension.git
cd /path/to/ComfyUI/custom_nodes
ln -s /path/to/ComfyUI-to-Python-Extension ComfyUI-to-Python-Extension
cd /path/to/ComfyUI-to-Python-Extension
uv sync
```
Then install this extension into the same Python environment that launches ComfyUI.
The `pyproject.toml` file declares the package dependencies, but those dependencies still need to be installed into ComfyUI's runtime Python.
If you run ComfyUI from a source checkout with `uv`:
```bash
cd /path/to/ComfyUI
uv pip install -e ./custom_nodes/ComfyUI-to-Python-Extension
uv run python main.py
```
If you use the Windows portable build:
```
cd C:\path\to\ComfyUI_windows_portable\ComfyUI\custom_nodes\ComfyUI-to-Python-Extension
..\..\..\python_embeded\python.exe -m pip install -e .
```
Running `uv sync` inside `ComfyUI-to-Python-Extension` creates this extension's own `.venv`.
ComfyUI does not automatically import dependencies from that `.venv`; it imports custom nodes with the Python interpreter used to launch ComfyUI.
After installation, restart ComfyUI.
### CLI exporter / generated scripts
@@ -55,9 +72,18 @@ export COMFYUI_PATH=/path/to/ComfyUI
```
`COMFYUI_PATH` helps the exporter and generated scripts find the ComfyUI codebase. It does not, by itself, register this repo as a ComfyUI extension for the Web UI.
It also does not install ComfyUI runtime dependencies such as `torch` into the current Python environment.
`COMFYUI_PATH` is checked first. If it is not set, the exporter falls back to searching parent directories for a folder named `ComfyUI`.
If the CLI fails with `ModuleNotFoundError: No module named 'torch'`, run the command with the same Python environment that launches ComfyUI, or install ComfyUI's runtime dependencies into the environment you are using for the CLI.
For Windows portable builds, run the CLI with ComfyUI's embedded Python from the extension directory:
```
..\..\..\python_embeded\python.exe -m comfyui_to_python --input_file ".\workflow_api.json" --output_file ".\workflow_api.py"
```
## Web UI Export
In current ComfyUI builds, `Save As Script` is typically available under:
@@ -65,12 +91,13 @@ In current ComfyUI builds, `Save As Script` is typically available under:
`File -> Save As Script`
The command downloads a generated `.py` file.
The current UI export uses the default filename `workflow_api.py` so it works in ComfyUI Desktop without relying on `prompt()`.
![Save As Script](images/save_as_script.png)
Notes:
- menu placement can differ between frontend versions
- ComfyUI Desktop may fail on the current filename prompt flow; use the CLI flow below if that happens
- the Web UI export uses a fixed default filename rather than asking for one interactively
## CLI Export
@@ -118,14 +145,30 @@ The generated script is a workflow export. It does not automatically turn workfl
Scripts exported directly from `File -> Save As Script` in the ComfyUI UI already include the frontend workflow metadata needed for drag-and-drop reimport. Images saved by those scripts can be dropped back into ComfyUI and reopen with the original workflow metadata.
Generated scripts reuse ComfyUI's runtime argument parser during bootstrap, so common ComfyUI memory flags such as `--highvram`, `--normalvram`, `--lowvram`, `--novram`, `--cpu`, and `--disable-smart-memory` can be passed directly to the exported `.py` file.
Lifecycle notes:
- exported scripts are single-shot workflow runners, not long-lived ComfyUI prompt servers
- they do not implement Web UI prompt/result caching across repeated service calls
- exported `main()` now performs best-effort ComfyUI model/cache cleanup in a `finally` block
- set `COMFYUI_TOPYTHON_UNLOAD_MODELS=1` or call `main(unload_models=True)` if an embedded or repeated-call host should aggressively unload models after each run instead of preserving them for reuse
## Troubleshooting
- unsupported Python version:
use Python 3.12 or newer in the environment that runs the extension, then reinstall the extension dependencies there
- Web UI import fails after `uv sync`:
`uv sync` in this repo installs dependencies into this repo's `.venv`, but ComfyUI loads custom nodes with its own Python environment. Install the extension into the Python interpreter that launches ComfyUI.
- Windows portable import fails after `uv sync`:
ComfyUI portable uses its bundled `python_embeded` interpreter. From the extension directory, run `..\..\..\python_embeded\python.exe -m pip install -e .`, then restart ComfyUI.
- CLI fails with `No module named 'torch'`:
the extension `.venv` may not have ComfyUI's runtime dependencies. Either run the CLI from the Python environment that launches ComfyUI, or make sure the target ComfyUI environment is installed and `COMFYUI_PATH` points to it.
- `Save As Script` not visible:
check your current ComfyUI menu/frontend version and look under `File`
- `Save As Script` not visible after restart:
make sure this repo is discoverable by ComfyUI through `custom_nodes` by cloning it into `ComfyUI/custom_nodes`, symlinking it there, or adding an external `custom_nodes` path in `extra_model_paths.yaml`
- Desktop says `prompt()` is unsupported:
use the CLI export flow instead
- save uses the default filename:
rename `workflow_api.py` after download if you want a different local filename
- ComfyUI cannot be found:
set `COMFYUI_PATH`
- models or paths are missing at runtime:
+4 -12
View File
@@ -13,18 +13,10 @@ sys.path.append(ext_dir)
try:
import black
except ImportError:
print("Unable to import requirements for ComfyUI-SaveAsScript.")
print("Installing...")
import importlib
spec = importlib.util.spec_from_file_location(
"impact_install", os.path.join(os.path.dirname(__file__), "install.py")
)
impact_install = importlib.util.module_from_spec(spec)
spec.loader.exec_module(impact_install)
print("Successfully installed. Hopefully, at least.")
raise ImportError(
"ComfyUI-to-Python-Extension requires the project dependencies to be installed. "
f"Run 'uv sync' in {ext_dir} with Python 3.12+ before loading this extension."
) from None
# Prevent reimporting of custom nodes
os.environ["RUNNING_IN_COMFYUI"] = "TRUE"
@@ -2,6 +2,7 @@ from ..node_runtime import (
add_comfyui_directory_to_sys_path,
add_extra_model_paths,
bootstrap_comfyui_runtime,
cleanup_comfyui_runtime,
find_path,
get_comfyui_path,
get_value_at_index,
@@ -11,6 +12,7 @@ __all__ = [
"add_comfyui_directory_to_sys_path",
"add_extra_model_paths",
"bootstrap_comfyui_runtime",
"cleanup_comfyui_runtime",
"find_path",
"get_comfyui_path",
"get_value_at_index",
+41 -8
View File
@@ -48,6 +48,7 @@ class WorkflowPlanner:
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()
input_value_types = self.get_input_value_types(input_types)
class_def = self.node_class_mappings[class_type]()
missing_required_variable = False
@@ -106,7 +107,7 @@ class WorkflowPlanner:
)
inputs = self.update_inputs(inputs, executed_variables)
seed_sync_code = self.create_prompt_seed_sync_code(
idx, inputs, is_special_function
idx, inputs, input_value_types, is_special_function
)
target_lines = special_functions_code if is_special_function else code
@@ -118,6 +119,7 @@ class WorkflowPlanner:
class_def.FUNCTION,
executed_variables[idx],
is_special_function,
input_value_types=input_value_types,
**inputs,
)
)
@@ -138,16 +140,24 @@ class WorkflowPlanner:
func: str,
variable_name: str,
is_special_function: bool,
input_value_types: dict[str, str] | None = None,
**kwargs,
) -> str:
args = ", ".join(self.format_arg(key, value) for key, value in kwargs.items())
args = ", ".join(
self.format_arg(key, value, (input_value_types or {}).get(key))
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
self,
node_id: str,
inputs: dict,
input_value_types: dict[str, str],
is_special_function: bool,
) -> list[str]:
seed_sync_lines = []
for key in ("seed", "noise_seed"):
@@ -156,8 +166,11 @@ class WorkflowPlanner:
randomized_seed_variable = (
f"node_{self.sanitize_node_id(str(node_id))}_{self.clean_variable_name(key)}"
)
randomized_seed_code = self.get_randomized_seed_code(
input_value_types.get(key)
)
seed_sync_lines.append(
f'{randomized_seed_variable} = prompt["{node_id}"]["inputs"]["{key}"] = random.randint(1, 2**64)'
f'{randomized_seed_variable} = prompt["{node_id}"]["inputs"]["{key}"] = {randomized_seed_code}'
)
inputs[key] = {"variable_name": randomized_seed_variable}
@@ -167,22 +180,42 @@ class WorkflowPlanner:
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)
def format_arg(self, key: str, value: Any, input_value_type: str | None = None) -> str:
value_code = self.format_arg_value(key, value, input_value_type)
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:
def format_arg_value(
key: str, value: Any, input_value_type: str | None = None
) -> 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)"
return WorkflowPlanner.get_randomized_seed_code(input_value_type)
if isinstance(value, str):
return json.dumps(value)
return repr(value)
@staticmethod
def get_input_value_types(input_types: dict) -> dict[str, str]:
value_types = {}
for section in ("required", "optional", "hidden"):
for key, value in input_types.get(section, {}).items():
if isinstance(value, tuple) and value:
value_types[key] = value[0]
elif isinstance(value, str):
value_types[key] = value
return value_types
@staticmethod
def get_randomized_seed_code(input_value_type: str | None) -> str:
randomized_seed_code = "random.randint(1, 2**64)"
if input_value_type == "STRING":
return f"str({randomized_seed_code})"
return randomized_seed_code
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"
+14 -5
View File
@@ -9,6 +9,7 @@ from .generated_helpers import (
add_comfyui_directory_to_sys_path,
add_extra_model_paths,
bootstrap_comfyui_runtime,
cleanup_comfyui_runtime,
find_path,
get_comfyui_path,
get_value_at_index,
@@ -36,6 +37,7 @@ class WorkflowRenderer:
add_comfyui_directory_to_sys_path,
add_extra_model_paths,
bootstrap_comfyui_runtime,
cleanup_comfyui_runtime,
]:
func_strings.append(f"\n{inspect.getsource(func)}")
@@ -74,7 +76,7 @@ class WorkflowRenderer:
execution_section = [
"# Workflow execution",
"def main():",
"def main(unload_models: bool | None = None):",
" bootstrap_comfyui_runtime()",
" add_extra_model_paths()",
]
@@ -88,20 +90,27 @@ class WorkflowRenderer:
"",
" import torch",
"",
" with torch.inference_mode():",
" try:",
" with torch.inference_mode():",
]
)
execution_section.extend(
self.build_function_body(
plan.special_functions_code, "pass", indentation=" "
plan.special_functions_code, "pass", indentation=" "
).splitlines()
)
execution_section.append(f" for q in range({plan.queue_size}):")
execution_section.append(f" for q in range({plan.queue_size}):")
execution_section.extend(
self.build_function_body(
plan.loop_code, "pass", indentation=" "
plan.loop_code, "pass", indentation=" "
).splitlines()
)
execution_section.extend(
[
" finally:",
" cleanup_comfyui_runtime(unload_models=unload_models)",
]
)
entrypoint_section = [
"# Entrypoint",
+39
View File
@@ -1,5 +1,6 @@
import os
import sys
import warnings
from typing import Any, Mapping, Sequence, Union
@@ -94,6 +95,44 @@ def bootstrap_comfyui_runtime() -> None:
os.environ["OCL_SET_SVM_SIZE"] = "262144"
def cleanup_comfyui_runtime(unload_models: bool | None = None) -> None:
"""Best-effort cleanup for embedded or repeated generated-script execution."""
import gc
def run_cleanup_hook(name: str, should_run: bool = True) -> None:
if not should_run or not hasattr(model_management, name):
return
cleanup_fn = getattr(model_management, name)
try:
cleanup_fn()
except Exception as exc:
warnings.warn(
f"ComfyUI cleanup hook {name} failed during teardown: {exc}",
RuntimeWarning,
stacklevel=2,
)
should_unload = unload_models
if should_unload is None:
should_unload = os.environ.get("COMFYUI_TOPYTHON_UNLOAD_MODELS", "").lower() in {
"1",
"true",
"yes",
"on",
}
try:
import comfy.model_management as model_management
except ModuleNotFoundError:
gc.collect()
return
run_cleanup_hook("cleanup_models_gc")
run_cleanup_hook("unload_all_models", should_run=should_unload)
run_cleanup_hook("soft_empty_cache")
gc.collect()
def import_custom_nodes() -> None:
"""Initialize ComfyUI custom nodes in the exporter runtime."""
comfyui_path = get_comfyui_path()
+6 -10
View File
@@ -1,6 +1,9 @@
import { api } from "../../scripts/api.js";
import { app } from "../../scripts/app.js";
const DEFAULT_SCRIPT_FILENAME = "workflow_api.py";
const DEFAULT_WORKFLOW_NAME = "workflow_api.json";
function $el(tag, options = {}) {
const element = document.createElement(tag);
const { parent, style, ...props } = options;
@@ -39,15 +42,12 @@ const extension = {
});
},
savePythonScript() {
var filename = prompt("Save script as:");
if(filename === undefined || filename === null || filename === "") {
return
}
const filename = DEFAULT_SCRIPT_FILENAME;
app.graphToPrompt().then(async (p) => {
const frontendWorkflow = p.workflow ?? app.graph.serialize();
const json = JSON.stringify({
name: filename + ".json",
name: DEFAULT_WORKFLOW_NAME,
workflow: JSON.stringify(p.output, null, 2),
frontend_workflow: JSON.stringify(frontendWorkflow, null, 2),
}, null, 2); // convert the data to a JSON string
@@ -55,10 +55,6 @@ const extension = {
if(response.status == 200) {
const blob = new Blob([await response.text()], {type: "text/python;charset=utf-8"});
const url = URL.createObjectURL(blob);
if(!filename.endsWith(".py")) {
filename += ".py";
}
const a = $el("a", {
href: url,
download: filename,
+2 -1
View File
@@ -1,8 +1,9 @@
[project]
name = "comfyui-to-python-extension"
description = "This custom node allows you to generate pure python code from your ComfyUI workflow with the click of a button. Great for rapid experimentation or production deployment."
version = "2.0.0"
version = "2.1.0"
license = { text = "MIT License" }
requires-python = ">=3.12"
dependencies = ["black"]
[project.urls]
@@ -0,0 +1,8 @@
{
"1": {
"class_type": "StringSeedNode",
"inputs": {
"seed": "seed-placeholder"
}
}
}
@@ -0,0 +1,27 @@
{
"33": {
"class_type": "VaeDecode",
"inputs": {
"samples": "latent-placeholder"
}
},
"42:0": {
"class_type": "UpscaleModelLoader",
"inputs": {
"model_name": "4x-ultrasharp.safetensors"
}
},
"42:1": {
"class_type": "ImageUpscaleWithModel",
"inputs": {
"upscale_model": [
"42:0",
0
],
"image": [
"33",
0
]
}
}
}
@@ -0,0 +1,10 @@
{
"1": {
"class_type": "TextConcatenateNode",
"inputs": {
"delimiter": "",
"clean_whitespace": "true",
"text_b": "\\"
}
}
}
@@ -0,0 +1,41 @@
{
"11": {
"class_type": "DualClipLoader",
"inputs": {
"clip_name": "clip.safetensors"
}
},
"633": {
"class_type": "AnySwitchRgthree",
"inputs": {
"model": "model-placeholder"
}
},
"631": {
"class_type": "PowerLoraLoaderRgthree",
"inputs": {
"PowerLoraLoaderHeaderWidget": {
"type": "PowerLoraLoaderHeaderWidget"
},
"lora_1": {
"on": false,
"lora": "lora.safetensors",
"strength": 1.2
},
"lora_2": {
"on": false,
"lora": "lora2.safetensors",
"strength": 0.7
},
"\u2795 Add Lora": "",
"model": [
"633",
0
],
"clip": [
"11",
0
]
}
}
}
@@ -0,0 +1,8 @@
{
"1": {
"class_type": "WindowsPathNode",
"inputs": {
"path": "C:\\ComfyUI\\models\\upscale_models\\RealESRGAN_x4plus.safetensors"
}
}
}
+86 -43
View File
@@ -105,6 +105,46 @@ def bootstrap_comfyui_runtime() -> None:
os.environ["OCL_SET_SVM_SIZE"] = "262144"
def cleanup_comfyui_runtime(unload_models: bool | None = None) -> None:
"""Best-effort cleanup for embedded or repeated generated-script execution."""
import gc
def run_cleanup_hook(name: str, should_run: bool = True) -> None:
if not should_run or not hasattr(model_management, name):
return
cleanup_fn = getattr(model_management, name)
try:
cleanup_fn()
except Exception as exc:
warnings.warn(
f"ComfyUI cleanup hook {name} failed during teardown: {exc}",
RuntimeWarning,
stacklevel=2,
)
should_unload = unload_models
if should_unload is None:
should_unload = os.environ.get(
"COMFYUI_TOPYTHON_UNLOAD_MODELS", ""
).lower() in {
"1",
"true",
"yes",
"on",
}
try:
import comfy.model_management as model_management
except ModuleNotFoundError:
gc.collect()
return
run_cleanup_hook("cleanup_models_gc")
run_cleanup_hook("unload_all_models", should_run=should_unload)
run_cleanup_hook("soft_empty_cache")
gc.collect()
# Workflow data
def build_workflow() -> dict[str, Any]:
return {
@@ -163,7 +203,7 @@ extra_pnginfo = build_extra_pnginfo()
# Workflow execution
def main():
def main(unload_models: bool | None = None):
bootstrap_comfyui_runtime()
add_extra_model_paths()
@@ -180,51 +220,54 @@ def main():
import torch
with torch.inference_mode():
checkpointloadersimple = CheckpointLoaderSimple()
checkpointloadersimple_1 = checkpointloadersimple.load_checkpoint(
ckpt_name="v1-5-pruned-emaonly-fp16.safetensors"
)
cliptextencode = CLIPTextEncode()
cliptextencode_2 = cliptextencode.encode(
text="a small cottage in a meadow, soft daylight",
clip=get_value_at_index(checkpointloadersimple_1, 1),
)
cliptextencode_3 = cliptextencode.encode(
text="blurry, low quality",
clip=get_value_at_index(checkpointloadersimple_1, 1),
)
emptylatentimage = EmptyLatentImage()
emptylatentimage_4 = emptylatentimage.generate(
width=512, height=512, batch_size=1
)
ksampler = KSampler()
vaedecode = VAEDecode()
saveimage = SaveImage()
for q in range(1):
node_5_seed = prompt["5"]["inputs"]["seed"] = random.randint(1, 2**64)
ksampler_5 = ksampler.sample(
seed=node_5_seed,
steps=4,
cfg=7,
sampler_name="euler",
scheduler="normal",
denoise=1,
model=get_value_at_index(checkpointloadersimple_1, 0),
positive=get_value_at_index(cliptextencode_2, 0),
negative=get_value_at_index(cliptextencode_3, 0),
latent_image=get_value_at_index(emptylatentimage_4, 0),
try:
with torch.inference_mode():
checkpointloadersimple = CheckpointLoaderSimple()
checkpointloadersimple_1 = checkpointloadersimple.load_checkpoint(
ckpt_name="v1-5-pruned-emaonly-fp16.safetensors"
)
vaedecode_6 = vaedecode.decode(
samples=get_value_at_index(ksampler_5, 0),
vae=get_value_at_index(checkpointloadersimple_1, 2),
cliptextencode = CLIPTextEncode()
cliptextencode_2 = cliptextencode.encode(
text="a small cottage in a meadow, soft daylight",
clip=get_value_at_index(checkpointloadersimple_1, 1),
)
saveimage_7 = saveimage.save_images(
filename_prefix="E2E_text_to_image",
images=get_value_at_index(vaedecode_6, 0),
prompt=prompt,
extra_pnginfo=extra_pnginfo,
cliptextencode_3 = cliptextencode.encode(
text="blurry, low quality",
clip=get_value_at_index(checkpointloadersimple_1, 1),
)
emptylatentimage = EmptyLatentImage()
emptylatentimage_4 = emptylatentimage.generate(
width=512, height=512, batch_size=1
)
ksampler = KSampler()
vaedecode = VAEDecode()
saveimage = SaveImage()
for q in range(1):
node_5_seed = prompt["5"]["inputs"]["seed"] = random.randint(1, 2**64)
ksampler_5 = ksampler.sample(
seed=node_5_seed,
steps=4,
cfg=7,
sampler_name="euler",
scheduler="normal",
denoise=1,
model=get_value_at_index(checkpointloadersimple_1, 0),
positive=get_value_at_index(cliptextencode_2, 0),
negative=get_value_at_index(cliptextencode_3, 0),
latent_image=get_value_at_index(emptylatentimage_4, 0),
)
vaedecode_6 = vaedecode.decode(
samples=get_value_at_index(ksampler_5, 0),
vae=get_value_at_index(checkpointloadersimple_1, 2),
)
saveimage_7 = saveimage.save_images(
filename_prefix="E2E_text_to_image",
images=get_value_at_index(vaedecode_6, 0),
prompt=prompt,
extra_pnginfo=extra_pnginfo,
)
finally:
cleanup_comfyui_runtime(unload_models=unload_models)
# Entrypoint
+53 -20
View File
@@ -105,6 +105,36 @@ def bootstrap_comfyui_runtime() -> None:
os.environ["OCL_SET_SVM_SIZE"] = "262144"
def cleanup_comfyui_runtime(unload_models: bool | None = None) -> None:
"""Best-effort cleanup for embedded or repeated generated-script execution."""
import gc
should_unload = unload_models
if should_unload is None:
should_unload = os.environ.get(
"COMFYUI_TOPYTHON_UNLOAD_MODELS", ""
).lower() in {
"1",
"true",
"yes",
"on",
}
try:
import comfy.model_management as model_management
except ModuleNotFoundError:
gc.collect()
return
if hasattr(model_management, "cleanup_models_gc"):
model_management.cleanup_models_gc()
if should_unload and hasattr(model_management, "unload_all_models"):
model_management.unload_all_models()
if hasattr(model_management, "soft_empty_cache"):
model_management.soft_empty_cache()
gc.collect()
def import_custom_nodes() -> None:
"""Initialize ComfyUI custom nodes in the exporter runtime."""
comfyui_path = get_comfyui_path()
@@ -160,7 +190,7 @@ extra_pnginfo = build_extra_pnginfo()
# Workflow execution
def main():
def main(unload_models: bool | None = None):
bootstrap_comfyui_runtime()
add_extra_model_paths()
import_custom_nodes()
@@ -170,26 +200,29 @@ def main():
import torch
with torch.inference_mode():
loadimage = LoadImage()
loadimage_1 = loadimage.load_image(image="e2e_upscale_input.png")
upscalemodelloader = NODE_CLASS_MAPPINGS["UpscaleModelLoader"]()
upscalemodelloader_2 = upscalemodelloader.EXECUTE_NORMALIZED(
model_name="RealESRGAN_x4plus.safetensors"
)
imageupscalewithmodel = NODE_CLASS_MAPPINGS["ImageUpscaleWithModel"]()
saveimage = SaveImage()
for q in range(1):
imageupscalewithmodel_3 = imageupscalewithmodel.EXECUTE_NORMALIZED(
upscale_model=get_value_at_index(upscalemodelloader_2, 0),
image=get_value_at_index(loadimage_1, 0),
)
saveimage_4 = saveimage.save_images(
filename_prefix="E2E_upscale_model_loader",
images=get_value_at_index(imageupscalewithmodel_3, 0),
prompt=prompt,
extra_pnginfo=extra_pnginfo,
try:
with torch.inference_mode():
loadimage = LoadImage()
loadimage_1 = loadimage.load_image(image="e2e_upscale_input.png")
upscalemodelloader = NODE_CLASS_MAPPINGS["UpscaleModelLoader"]()
upscalemodelloader_2 = upscalemodelloader.EXECUTE_NORMALIZED(
model_name="RealESRGAN_x4plus.safetensors"
)
imageupscalewithmodel = NODE_CLASS_MAPPINGS["ImageUpscaleWithModel"]()
saveimage = SaveImage()
for q in range(1):
imageupscalewithmodel_3 = imageupscalewithmodel.EXECUTE_NORMALIZED(
upscale_model=get_value_at_index(upscalemodelloader_2, 0),
image=get_value_at_index(loadimage_1, 0),
)
saveimage_4 = saveimage.save_images(
filename_prefix="E2E_upscale_model_loader",
images=get_value_at_index(imageupscalewithmodel_3, 0),
prompt=prompt,
extra_pnginfo=extra_pnginfo,
)
finally:
cleanup_comfyui_runtime(unload_models=unload_models)
# Entrypoint
@@ -0,0 +1,312 @@
import json
import unittest
from io import StringIO
from pathlib import Path
from comfyui_to_python import ComfyUItoPython
FIXTURE_DIR = Path(__file__).parent / "fixtures" / "unit" / "generator_codegen"
class AnySwitchRgthree:
CATEGORY = "utils"
FUNCTION = "switch"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("MODEL",),
}
}
def switch(self, model):
return (model,)
class DualClipLoader:
CATEGORY = "loaders"
FUNCTION = "load_clip"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"clip_name": ("STRING",),
}
}
def load_clip(self, clip_name):
return (clip_name,)
class PowerLoraLoaderRgthree:
CATEGORY = "loaders"
FUNCTION = "load_loras"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"PowerLoraLoaderHeaderWidget": ("DICT",),
"model": ("MODEL",),
"clip": ("CLIP",),
}
}
def load_loras(self, **kwargs):
return (kwargs,)
class UpscaleModelLoader:
CATEGORY = "loaders"
FUNCTION = "load_model"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model_name": ("STRING",),
}
}
def load_model(self, model_name):
return (model_name,)
class ImageUpscaleWithModel:
CATEGORY = "image/upscaling"
FUNCTION = "upscale"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"upscale_model": ("UPSCALE_MODEL",),
"image": ("IMAGE",),
}
}
def upscale(self, upscale_model, image):
return (image,)
class VaeDecode:
CATEGORY = "latent"
FUNCTION = "decode"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"samples": ("LATENT",),
}
}
def decode(self, samples):
return (samples,)
class WindowsPathNode:
CATEGORY = "paths"
FUNCTION = "open_path"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"path": ("STRING",),
}
}
def open_path(self, path):
return (path,)
class TextConcatenateNode:
CATEGORY = "text"
FUNCTION = "text_concatenate"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"delimiter": ("STRING",),
"clean_whitespace": ("STRING",),
"text_b": ("STRING",),
}
}
def text_concatenate(self, delimiter, clean_whitespace, text_b):
return (delimiter, clean_whitespace, text_b)
class StringSeedNode:
CATEGORY = "sampling"
FUNCTION = "sample"
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"seed": ("STRING",),
},
"hidden": {
"prompt": "PROMPT",
},
}
def sample(self, seed, prompt):
return (seed, prompt)
def load_fixture(name: str) -> dict:
return json.loads((FIXTURE_DIR / name).read_text(encoding="utf-8"))
def export_workflow(workflow: dict, node_class_mappings: dict) -> str:
output = StringIO()
ComfyUItoPython(
workflow=json.dumps(workflow),
output_file=output,
node_class_mappings=node_class_mappings,
)
return output.getvalue()
class GeneratorCodegenIssueRegressionTest(unittest.TestCase):
def test_export_uses_dictionary_expansion_for_rgthree_symbol_heavy_input_names(self):
generated = export_workflow(
load_fixture("unsafe-rgthree-kwargs.json"),
{
"AnySwitchRgthree": AnySwitchRgthree,
"DualClipLoader": DualClipLoader,
"PowerLoraLoaderRgthree": PowerLoraLoaderRgthree,
},
)
self.assertIn(
'powerloraloaderrgthree_631 = powerloraloaderrgthree.load_loras(',
generated,
)
self.assertIn(
'PowerLoraLoaderHeaderWidget={"type": "PowerLoraLoaderHeaderWidget"}',
generated,
)
self.assertIn('**{"\\u2795 Add Lora": ""}', generated)
self.assertNotIn('➕ Add Lora=""', generated)
def test_export_sanitizes_subgraph_identifiers_for_upscaler_workflows(self):
generated = export_workflow(
load_fixture("subgraph-upscaler-identifiers.json"),
{
"VaeDecode": VaeDecode,
"UpscaleModelLoader": UpscaleModelLoader,
"ImageUpscaleWithModel": ImageUpscaleWithModel,
},
)
self.assertIn("upscalemodelloader_42_0 = upscalemodelloader.load_model(", generated)
self.assertIn(
"imageupscalewithmodel_42_1 = imageupscalewithmodel.upscale(",
generated,
)
self.assertIn(
"upscale_model=get_value_at_index(upscalemodelloader_42_0, 0)",
generated,
)
self.assertNotIn("imageupscalewithmodel_42:1", generated)
self.assertNotIn("upscalemodelloader_42:0", generated)
def test_export_preserves_windows_style_model_paths(self):
generated = export_workflow(
load_fixture("windows-path-string.json"),
{
"WindowsPathNode": WindowsPathNode,
},
)
globals_dict = {"__name__": "generated_workflow_module"}
exec(generated, globals_dict)
self.assertEqual(
globals_dict["build_workflow"]()["1"]["inputs"]["path"],
r"C:\ComfyUI\models\upscale_models\RealESRGAN_x4plus.safetensors",
)
def test_export_preserves_trailing_backslash_string_literals(self):
generated = export_workflow(
load_fixture("trailing-backslash-string.json"),
{
"TextConcatenateNode": TextConcatenateNode,
},
)
globals_dict = {"__name__": "generated_workflow_module"}
exec(generated, globals_dict)
self.assertEqual(
globals_dict["build_workflow"]()["1"]["inputs"]["text_b"],
"\\",
)
def test_export_randomizes_string_seed_inputs_as_strings(self):
generated = export_workflow(
load_fixture("string-seed-node.json"),
{
"StringSeedNode": StringSeedNode,
},
)
self.assertIn(
'node_1_seed = prompt["1"]["inputs"]["seed"] = str(random.randint(1, 2**64))',
generated,
)
self.assertIn("seed=node_1_seed", generated)
def test_issue_cluster_regressions_render_parseable_python(self):
workflows = [
(
load_fixture("unsafe-rgthree-kwargs.json"),
{
"AnySwitchRgthree": AnySwitchRgthree,
"DualClipLoader": DualClipLoader,
"PowerLoraLoaderRgthree": PowerLoraLoaderRgthree,
},
),
(
load_fixture("subgraph-upscaler-identifiers.json"),
{
"VaeDecode": VaeDecode,
"UpscaleModelLoader": UpscaleModelLoader,
"ImageUpscaleWithModel": ImageUpscaleWithModel,
},
),
(
load_fixture("trailing-backslash-string.json"),
{
"TextConcatenateNode": TextConcatenateNode,
},
),
(
load_fixture("windows-path-string.json"),
{
"WindowsPathNode": WindowsPathNode,
},
),
(
load_fixture("string-seed-node.json"),
{
"StringSeedNode": StringSeedNode,
},
),
]
for workflow, mapping in workflows:
generated = export_workflow(workflow, mapping)
compile(generated, "<generated_workflow>", "exec")
if __name__ == "__main__":
unittest.main()
+116
View File
@@ -0,0 +1,116 @@
import sys
import types
import unittest
import warnings
from unittest.mock import Mock, patch
from comfyui_to_python.node_runtime import cleanup_comfyui_runtime
class NodeRuntimeCleanupTest(unittest.TestCase):
def test_cleanup_releases_caches_without_forcing_model_unload_by_default(self):
comfy_module = types.ModuleType("comfy")
comfy_module.__path__ = []
model_management = types.ModuleType("comfy.model_management")
model_management.cleanup_models_gc = Mock()
model_management.unload_all_models = Mock()
model_management.soft_empty_cache = Mock()
comfy_module.model_management = model_management
with patch.dict(
sys.modules,
{
"comfy": comfy_module,
"comfy.model_management": model_management,
},
), patch.dict("os.environ", {}, clear=False):
cleanup_comfyui_runtime()
model_management.cleanup_models_gc.assert_called_once_with()
model_management.soft_empty_cache.assert_called_once_with()
model_management.unload_all_models.assert_not_called()
def test_cleanup_can_force_model_unload_from_environment(self):
comfy_module = types.ModuleType("comfy")
comfy_module.__path__ = []
model_management = types.ModuleType("comfy.model_management")
model_management.cleanup_models_gc = Mock()
model_management.unload_all_models = Mock()
model_management.soft_empty_cache = Mock()
comfy_module.model_management = model_management
with patch.dict(
sys.modules,
{
"comfy": comfy_module,
"comfy.model_management": model_management,
},
), patch.dict(
"os.environ",
{"COMFYUI_TOPYTHON_UNLOAD_MODELS": "true"},
clear=False,
):
cleanup_comfyui_runtime()
model_management.unload_all_models.assert_called_once_with()
def test_cleanup_suppresses_hook_failures_and_warns(self):
comfy_module = types.ModuleType("comfy")
comfy_module.__path__ = []
model_management = types.ModuleType("comfy.model_management")
model_management.cleanup_models_gc = Mock(side_effect=RuntimeError("gc failed"))
model_management.unload_all_models = Mock(side_effect=RuntimeError("unload failed"))
model_management.soft_empty_cache = Mock(side_effect=RuntimeError("cache failed"))
comfy_module.model_management = model_management
with patch.dict(
sys.modules,
{
"comfy": comfy_module,
"comfy.model_management": model_management,
},
), warnings.catch_warnings(record=True) as caught:
warnings.simplefilter("always")
cleanup_comfyui_runtime(unload_models=True)
model_management.cleanup_models_gc.assert_called_once_with()
model_management.unload_all_models.assert_called_once_with()
model_management.soft_empty_cache.assert_called_once_with()
self.assertEqual(len(caught), 3)
self.assertEqual(
[str(warning.message) for warning in caught],
[
"ComfyUI cleanup hook cleanup_models_gc failed during teardown: gc failed",
"ComfyUI cleanup hook unload_all_models failed during teardown: unload failed",
"ComfyUI cleanup hook soft_empty_cache failed during teardown: cache failed",
],
)
def test_cleanup_does_not_mask_active_workflow_exception(self):
comfy_module = types.ModuleType("comfy")
comfy_module.__path__ = []
model_management = types.ModuleType("comfy.model_management")
model_management.cleanup_models_gc = Mock(side_effect=RuntimeError("cleanup failed"))
model_management.soft_empty_cache = Mock()
comfy_module.model_management = model_management
with patch.dict(
sys.modules,
{
"comfy": comfy_module,
"comfy.model_management": model_management,
},
), warnings.catch_warnings(record=True):
warnings.simplefilter("always")
with self.assertRaisesRegex(ValueError, "workflow failed"):
try:
raise ValueError("workflow failed")
finally:
cleanup_comfyui_runtime()
model_management.cleanup_models_gc.assert_called_once_with()
model_management.soft_empty_cache.assert_called_once_with()
if __name__ == "__main__":
unittest.main()
+33
View File
@@ -0,0 +1,33 @@
import tomllib
import unittest
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parent.parent
class ProjectContractsTest(unittest.TestCase):
def test_project_declares_supported_python_floor(self):
pyproject = tomllib.loads((REPO_ROOT / "pyproject.toml").read_text(encoding="utf-8"))
self.assertEqual(pyproject["project"]["requires-python"], ">=3.12")
def test_readme_documents_python_support_and_default_save_filename(self):
readme = (REPO_ROOT / "README.md").read_text(encoding="utf-8")
self.assertIn("This project supports Python 3.12 and newer.", readme)
self.assertIn("default filename `workflow_api.py`", readme)
def test_extension_import_path_requires_uv_sync_instead_of_running_install_py(self):
init_text = (REPO_ROOT / "__init__.py").read_text(encoding="utf-8")
self.assertIn("Run 'uv sync'", init_text)
self.assertNotIn("spec_from_file_location", init_text)
self.assertNotIn("Successfully installed. Hopefully, at least.", init_text)
def test_frontend_save_flow_uses_deterministic_filename_without_prompt(self):
save_as_script = (REPO_ROOT / "js" / "save-as-script.js").read_text(encoding="utf-8")
self.assertIn('const DEFAULT_SCRIPT_FILENAME = "workflow_api.py";', save_as_script)
self.assertNotIn("prompt(", save_as_script)
+18 -6
View File
@@ -145,20 +145,26 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
generated = output.getvalue()
self.assertIn("def bootstrap_comfyui_runtime()", generated)
self.assertIn("def cleanup_comfyui_runtime(", generated)
self.assertIn("import comfy.options", generated)
self.assertIn("comfy.options.enable_args_parsing()", generated)
self.assertIn("import cuda_malloc", generated)
self.assertNotIn("\nbootstrap_comfyui_runtime()\n", generated)
self.assertIn(
"def main():\n bootstrap_comfyui_runtime()\n add_extra_model_paths()",
"def main(unload_models: bool | None = None):\n"
" bootstrap_comfyui_runtime()\n"
" add_extra_model_paths()",
generated,
)
self.assertLess(
generated.index("def bootstrap_comfyui_runtime()"), generated.index("def main():")
generated.index("def bootstrap_comfyui_runtime()"),
generated.index("def main(unload_models: bool | None = None):"),
)
main_section = generated[generated.index("def main():") :]
main_section = generated[
generated.index("def main(unload_models: bool | None = None):") :
]
self.assertIn(
"def main():\n"
"def main(unload_models: bool | None = None):\n"
" bootstrap_comfyui_runtime()\n"
" add_extra_model_paths()",
main_section,
@@ -172,6 +178,10 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
main_section.index("import torch"),
)
self.assertLess(generated.index("import cuda_malloc"), generated.index("import torch"))
self.assertIn(
" finally:\n cleanup_comfyui_runtime(unload_models=unload_models)",
main_section,
)
def test_generated_module_import_does_not_parse_cli_args(self):
workflow = {
@@ -431,18 +441,20 @@ class UpscaleModelLoaderExportTest(unittest.TestCase):
self.assertIn("# Entrypoint", generated)
self.assertIn("def build_workflow()", generated)
self.assertIn("def build_extra_pnginfo()", generated)
self.assertIn("def main()", generated)
self.assertIn("def main(unload_models: bool | None = None)", generated)
self.assertIn("bootstrap_comfyui_runtime()", generated)
self.assertIn("cleanup_comfyui_runtime(unload_models=unload_models)", generated)
self.assertNotIn("def initialize_workflow()", generated)
self.assertNotIn("def run_once(", generated)
self.assertIn("with torch.inference_mode():", generated)
self.assertIn("finally:", generated)
self.assertIn("for q in range(1):", generated)
self.assertIn("workflow = build_workflow()", generated)
self.assertIn("extra_pnginfo = build_extra_pnginfo()", generated)
self.assertNotIn('workflow = json.loads("', generated)
self.assertLess(
generated.index("def build_workflow()"),
generated.index("def main()"),
generated.index("def main(unload_models: bool | None = None)"),
)
def test_hidden_metadata_kwargs_follow_function_signature(self):
Generated
+1 -1
View File
@@ -57,7 +57,7 @@ wheels = [
[[package]]
name = "comfyui-to-python-extension"
version = "2.0.0"
version = "2.1.0"
source = { virtual = "." }
dependencies = [
{ name = "black" },