diff --git a/README.md b/README.md index d4ce020..274ad8a 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,7 @@ 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`) @@ -65,12 +66,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 +120,24 @@ 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, then rerun `uv sync` - `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: diff --git a/comfyui_to_python/generator/generated_helpers.py b/comfyui_to_python/generator/generated_helpers.py index 88f5f2f..896ec37 100644 --- a/comfyui_to_python/generator/generated_helpers.py +++ b/comfyui_to_python/generator/generated_helpers.py @@ -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", diff --git a/comfyui_to_python/generator/render.py b/comfyui_to_python/generator/render.py index 1c531bc..1acd4a3 100644 --- a/comfyui_to_python/generator/render.py +++ b/comfyui_to_python/generator/render.py @@ -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", diff --git a/comfyui_to_python/node_runtime.py b/comfyui_to_python/node_runtime.py index 894e317..17028c8 100644 --- a/comfyui_to_python/node_runtime.py +++ b/comfyui_to_python/node_runtime.py @@ -94,6 +94,34 @@ 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() diff --git a/tests/test_node_runtime_cleanup.py b/tests/test_node_runtime_cleanup.py new file mode 100644 index 0000000..2054ef1 --- /dev/null +++ b/tests/test_node_runtime_cleanup.py @@ -0,0 +1,58 @@ +import sys +import types +import unittest +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() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_upscale_model_loader_export.py b/tests/test_upscale_model_loader_export.py index 804de12..9c1ff38 100644 --- a/tests/test_upscale_model_loader_export.py +++ b/tests/test_upscale_model_loader_export.py @@ -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):