Guard ComfyUI cleanup teardown failures

Make generator codegen regressions explicit and fix string seed rand

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:

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",

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",

comfyui_to_python/node_runtime.py

@@ -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()

tests/runtime/generated/text-to-image.py

@@ -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

tests/runtime/generated/upscale-model-loader.py

@@ -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

tests/test_node_runtime_cleanup.py

@@ -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()

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):