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