From f9163f23752a037b8fc9c37c308a0304969efd6b Mon Sep 17 00:00:00 2001 From: Peyton DeNiro <25550995+pydn@users.noreply.github.com> Date: Sun, 29 Mar 2026 18:18:37 -0500 Subject: [PATCH 1/2] 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. --- README.md | 18 +++++- .../generator/generated_helpers.py | 2 + comfyui_to_python/generator/render.py | 19 ++++-- comfyui_to_python/node_runtime.py | 28 +++++++++ tests/test_node_runtime_cleanup.py | 58 +++++++++++++++++++ tests/test_upscale_model_loader_export.py | 24 ++++++-- 6 files changed, 135 insertions(+), 14 deletions(-) create mode 100644 tests/test_node_runtime_cleanup.py 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): From 77c0bbec9bf70657427852b07776a53437e90a6b Mon Sep 17 00:00:00 2001 From: Peyton DeNiro <25550995+pydn@users.noreply.github.com> Date: Sun, 29 Mar 2026 19:02:21 -0500 Subject: [PATCH 2/2] Guard ComfyUI cleanup teardown failures --- comfyui_to_python/node_runtime.py | 23 +++- tests/runtime/generated/text-to-image.py | 129 ++++++++++++------ .../runtime/generated/upscale-model-loader.py | 73 +++++++--- tests/test_node_runtime_cleanup.py | 58 ++++++++ 4 files changed, 214 insertions(+), 69 deletions(-) diff --git a/comfyui_to_python/node_runtime.py b/comfyui_to_python/node_runtime.py index 17028c8..a74e9b0 100644 --- a/comfyui_to_python/node_runtime.py +++ b/comfyui_to_python/node_runtime.py @@ -1,5 +1,6 @@ import os import sys +import warnings from typing import Any, Mapping, Sequence, Union @@ -98,6 +99,19 @@ 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 { @@ -113,12 +127,9 @@ def cleanup_comfyui_runtime(unload_models: bool | None = None) -> None: 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() + run_cleanup_hook("cleanup_models_gc") + run_cleanup_hook("unload_all_models", should_run=should_unload) + run_cleanup_hook("soft_empty_cache") gc.collect() diff --git a/tests/runtime/generated/text-to-image.py b/tests/runtime/generated/text-to-image.py index e57de44..99c66dc 100644 --- a/tests/runtime/generated/text-to-image.py +++ b/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 diff --git a/tests/runtime/generated/upscale-model-loader.py b/tests/runtime/generated/upscale-model-loader.py index ed56d7a..98ed0ae 100644 --- a/tests/runtime/generated/upscale-model-loader.py +++ b/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 diff --git a/tests/test_node_runtime_cleanup.py b/tests/test_node_runtime_cleanup.py index 2054ef1..3379ef5 100644 --- a/tests/test_node_runtime_cleanup.py +++ b/tests/test_node_runtime_cleanup.py @@ -1,6 +1,7 @@ import sys import types import unittest +import warnings from unittest.mock import Mock, patch from comfyui_to_python.node_runtime import cleanup_comfyui_runtime @@ -53,6 +54,63 @@ class NodeRuntimeCleanupTest(unittest.TestCase): 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()