Upload salia_maxVRAMpurge.py

salia_maxVRAMpurge.py

@@ -0,0 +1,233 @@
+from __future__ import annotations
+
+import fnmatch
+import gc
+from typing import Any
+
+import torch
+import comfy.model_management as mm
+
+
+class AnyType(str):
+    """Wildcard type helper for ComfyUI type checks."""
+
+    def __ne__(self, __value: object) -> bool:
+        return False
+
+
+any_type = AnyType("*")
+
+
+# -----------------------------------------------------------------------------
+# Verified behavior from LayerStyle: Purge VRAM V2
+# purge_vram_v2(...):
+#   clear_memory()
+#   unload_all_models()
+#   soft_empty_cache()
+#   return (anything,)
+#
+# Verified clear_memory() from LayerStyle/imagefunc.py:
+#   gc.collect()
+#   if torch.cuda.is_available():
+#       torch.cuda.empty_cache()
+#       torch.cuda.ipc_collect()
+# -----------------------------------------------------------------------------
+def _layerstyle_clear_memory() -> None:
+    gc.collect()
+    try:
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+    except Exception:
+        pass
+
+
+def _run_layerstyle_purge_v2() -> None:
+    _layerstyle_clear_memory()
+
+    if hasattr(mm, "unload_all_models"):
+        try:
+            mm.unload_all_models()
+        except Exception:
+            pass
+
+    if hasattr(mm, "soft_empty_cache"):
+        try:
+            mm.soft_empty_cache()
+        except TypeError:
+            try:
+                mm.soft_empty_cache(False)
+            except Exception:
+                pass
+        except Exception:
+            pass
+
+
+# -----------------------------------------------------------------------------
+# Verified behavior from comfyui-unload-model: UnloadAllModels
+# route(...):
+#   unload_all_models()
+#   soft_empty_cache(True)
+#   gc.collect()
+#   torch.cuda.empty_cache()
+#   torch.cuda.ipc_collect()
+#   return (value,)
+# -----------------------------------------------------------------------------
+def _run_unload_all_models_pack() -> None:
+    if hasattr(mm, "unload_all_models"):
+        try:
+            mm.unload_all_models()
+        except Exception:
+            pass
+
+    if hasattr(mm, "soft_empty_cache"):
+        try:
+            mm.soft_empty_cache(True)
+        except TypeError:
+            try:
+                mm.soft_empty_cache(force=True)
+            except Exception:
+                try:
+                    mm.soft_empty_cache()
+                except Exception:
+                    pass
+        except Exception:
+            pass
+
+    gc.collect()
+
+    try:
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+    except Exception:
+        pass
+
+
+# -----------------------------------------------------------------------------
+# Missing behavior from the provided easy cleanGpuUsed replacement:
+#   remove_cache('*')
+# This is NOT present in the two researched nodes, so we add it here.
+# -----------------------------------------------------------------------------
+def remove_cache(cache_key: str = "*") -> None:
+    """Best-effort backend cache clearing for varying ComfyUI versions."""
+    pattern = cache_key or "*"
+
+    def _clear_dict(d: dict) -> None:
+        if pattern == "*":
+            d.clear()
+            return
+        for k in list(d.keys()):
+            if fnmatch.fnmatch(str(k), pattern):
+                d.pop(k, None)
+
+    def _clear_known_cache_dicts(obj: Any) -> None:
+        if obj is None:
+            return
+
+        interesting_names = {
+            "cache",
+            "caches",
+            "output_cache",
+            "outputs",
+            "outputs_ui",
+            "history",
+            "prompt_cache",
+        }
+
+        for name in dir(obj):
+            if name.startswith("_"):
+                continue
+
+            lname = name.lower()
+            if lname in {"queue", "task_queue", "tasks", "prompt_queue"}:
+                continue
+
+            if (lname in interesting_names) or ("cache" in lname):
+                try:
+                    val = getattr(obj, name, None)
+                except Exception:
+                    continue
+                if isinstance(val, dict):
+                    _clear_dict(val)
+
+    try:
+        import server
+
+        ps_cls = getattr(server, "PromptServer", None)
+        ps_inst = getattr(ps_cls, "instance", None) if ps_cls is not None else None
+        ps = ps_inst() if callable(ps_inst) else ps_inst
+
+        pq = getattr(ps, "prompt_queue", None) if ps is not None else None
+        _clear_known_cache_dicts(pq)
+        _clear_known_cache_dicts(getattr(pq, "executor", None))
+        _clear_known_cache_dicts(getattr(ps, "executor", None))
+        _clear_known_cache_dicts(getattr(ps, "prompt_executor", None))
+    except Exception:
+        pass
+
+    try:
+        import comfy_execution.caching as caching  # type: ignore
+
+        _clear_known_cache_dicts(caching)
+
+        for attr in ("CACHE", "CACHES", "cache", "caches"):
+            d = getattr(caching, attr, None)
+            if isinstance(d, dict):
+                _clear_dict(d)
+
+        for func_name in ("clear_cache", "clear_caches", "clear_all", "clear"):
+            fn = getattr(caching, func_name, None)
+            if callable(fn):
+                try:
+                    try:
+                        fn(pattern)
+                    except TypeError:
+                        fn()
+                except Exception:
+                    pass
+    except Exception:
+        pass
+
+    gc.collect()
+
+
+class Salia_Maximum_VRAM_Purge:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {"anything": (any_type, {})},
+            "optional": {},
+            "hidden": {"unique_id": "UNIQUE_ID", "extra_pnginfo": "EXTRA_PNGINFO"},
+        }
+
+    RETURN_TYPES = (any_type,)
+    RETURN_NAMES = ("any",)
+    OUTPUT_NODE = True
+    FUNCTION = "maximum_purge"
+    CATEGORY = "System/Memory"
+
+    @classmethod
+    def IS_CHANGED(cls, **kwargs):
+        return float("nan")
+
+    def maximum_purge(self, anything, unique_id=None, extra_pnginfo=None):
+        # 1) Run the verified LayerStyle Purge VRAM V2 behavior.
+        _run_layerstyle_purge_v2()
+
+        # 2) Run the verified comfyui-unload-model UnloadAllModels behavior.
+        _run_unload_all_models_pack()
+
+        # 3) Add only the missing easy cleanGpuUsed part that the two nodes do not do.
+        remove_cache("*")
+
+        return (anything,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "Salia_Maximum_VRAM_Purge": Salia_Maximum_VRAM_Purge,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "Salia_Maximum_VRAM_Purge": "Salia_Maximum_VRAM_Purge",
+}