Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6cdcc235a0 | ||
|
|
9a17000417 | ||
|
|
fbc06463f7 | ||
|
|
77c0bbec9b | ||
|
|
f9163f2375 | ||
|
|
d6c5b110ab | ||
|
|
5d70efdc88 | ||
|
|
228c3a076a | ||
|
|
574d96369c |
@@ -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()`.
|
||||
|
||||

|
||||
|
||||
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
@@ -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",
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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
@@ -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
@@ -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"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user