Guard ComfyUI cleanup teardown failures

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