nodes/pydn-ComfyUI-to-Python-Extension
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #153 from pydn/runtime-parity-cache-memory-lifecycle

Browse Source 
Add generated runtime cleanup after workflow execution
This commit is contained in:
Peyton DeNiro
2026-03-29 19:02:51 -05:00
committed by GitHub
parent d6c5b110ab 77c0bbec9b
commit fbc06463f7
8 changed files with 343 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/runtime/generated/text-to-image.py
+86 -43
View File
@@ -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
+53 -20
View File
@@ -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
+116
View File
@@ -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
+18 -6
View File
@@ -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):