feat: add Kiko Purge VRAM node for intelligent memory management

- Add comprehensive VRAM management tool with 4 purge modes (soft, aggressive, models_only, cache_only)
- Implement smart memory threshold triggering to avoid unnecessary purges
- Add detailed memory reporting showing before/after stats and freed MB
- Support passthrough design for seamless workflow integration
- Include graceful CPU fallback for non-CUDA environments
- Add comprehensive test suite with 14 tests covering all functionality
- Update documentation with detailed usage examples and parameters
- Create reusable AnyType class for wildcard input matching

The node provides essential memory management capabilities for complex workflows,
helping prevent OOM errors and optimize multi-model processing pipelines.

README.md

@@ -36,6 +36,7 @@ I’m sharing them here with the community, and I hope you find them as useful a
 | [📉 Image Scale Down By](#-image-scale-down-by) | Scale images down by a factor with quality preservation | 🖼️ Resolution |
 | [🎬 Film Grain](#-film-grain) | Add realistic film grain effects to images | 💾 Images |
 | [🔤 Embedding Autocomplete](#-embedding-autocomplete) | Smart autocomplete for embeddings, LoRAs, and tags | 🔧 Utils |
+| [🧹 Kiko Purge VRAM](#-kiko-purge-vram) | Intelligent VRAM management with detailed reporting | 🛠️ Utils |
 
 ### 🧰 xyz-helpers Tools
 
@@ -395,6 +396,46 @@ This feature is an enhanced fork of the autocomplete functionality from [ComfyUI
 - Original autocomplete concept by [pythongosssss](https://github.com/pythongosssss/ComfyUI-Custom-Scripts)
 - Enhanced and modernized by KikoTools team
 
+### 🧹 Kiko Purge VRAM
+**Intelligent GPU memory management with threshold-based triggering and detailed reporting.**
+
+**Key Features:**
+- **4 Purge Modes**: 
+  - `soft`: Basic garbage collection and cache clearing
+  - `aggressive`: Multiple GC passes with full CUDA cache clearing
+  - `models_only`: Unload all models and clear model cache
+  - `cache_only`: Clear CUDA cache without garbage collection
+- **Smart Thresholds**: Only purge when memory usage exceeds specified MB limit
+- **Detailed Reporting**: Shows before/after memory usage, freed MB, and timing
+- **Passthrough Design**: Acts as workflow checkpoint without disrupting data flow
+- **CPU Fallback**: Gracefully handles non-CUDA environments
+
+**Use Cases:**
+- Free memory between heavy processing stages
+- Prevent OOM errors in complex workflows
+- Debug memory usage patterns
+- Optimize multi-model workflows
+- Clean up after batch processing
+
+**Parameters:**
+- **anything**: Any input (passed through unchanged)
+- **mode**: Purge strategy selection
+- **report_memory**: Generate detailed memory statistics
+- **memory_threshold_mb**: Only purge if usage exceeds (0 = always purge)
+
+**Example Output:**
+```
+Memory usage (5000.0 MB) exceeds threshold (4000 MB)
+
+Memory Purge Report
+-------------------
+Mode: soft
+Memory Freed: 2500.0 MB
+Before: 5000.0 MB used (62.5%)
+After: 2500.0 MB used (31.3%)
+Time: 150.0ms
+```
+
 ### 💾 Kiko Save Image Features
 
 **Use Cases:**

kikotools/__init__.py

@@ -11,6 +11,7 @@ from .tools.gemini_prompt import GeminiPromptNode
 from .tools.image_scale_down_by import ImageScaleDownByNode
 from .tools.image_to_multiple_of import ImageToMultipleOfNode
 from .tools.kiko_film_grain import KikoFilmGrainNode
+from .tools.kiko_purge_vram import KikoPurgeVRAM
 from .tools.kiko_save_image import KikoSaveImageNode
 from .tools.resolution_calculator import ResolutionCalculatorNode
 from .tools.sampler_combo import SamplerComboCompactNode, SamplerComboNode
@@ -36,6 +37,7 @@ NODE_CLASS_MAPPINGS = {
     "DisplayAny": DisplayAnyNode,
     "DisplayText": DisplayTextNode,
     "KikoFilmGrain": KikoFilmGrainNode,
+    "KikoPurgeVRAM": KikoPurgeVRAM,
     "SamplerSelectHelper": SamplerSelectHelperNode,
     "SchedulerSelectHelper": SchedulerSelectHelperNode,
     "TextEncodeSamplerParams": TextEncodeSamplerParamsNode,
@@ -60,6 +62,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "DisplayAny": "Display Any",
     "DisplayText": "Display Text",
     "KikoFilmGrain": "Film Grain",
+    "KikoPurgeVRAM": "Kiko Purge VRAM",
     "SamplerSelectHelper": "Sampler Select Helper",
     "SchedulerSelectHelper": "Scheduler Select Helper",
     "TextEncodeSamplerParams": "Text Encode for Sampler Params",

kikotools/base/any_type.py

@@ -0,0 +1,8 @@
+"""AnyType for wildcard input matching in ComfyUI nodes."""
+
+
+class AnyType(str):
+    """A special type that matches any input type in ComfyUI."""
+
+    def __ne__(self, other):
+        return False

kikotools/tools/kiko_purge_vram/__init__.py

@@ -0,0 +1,3 @@
+from .node import KikoPurgeVRAM
+
+__all__ = ["KikoPurgeVRAM"]

kikotools/tools/kiko_purge_vram/logic.py

@@ -0,0 +1,130 @@
+import gc
+from typing import Dict, Tuple
+
+try:
+    import torch
+
+    TORCH_AVAILABLE = True
+except ImportError:
+    TORCH_AVAILABLE = False
+
+try:
+    import comfy.model_management as mm
+
+    COMFY_AVAILABLE = True
+except ImportError:
+    COMFY_AVAILABLE = False
+
+
+def get_memory_stats() -> Dict[str, float]:
+    stats = {
+        "cuda_available": False,
+        "free_mb": 0,
+        "total_mb": 0,
+        "used_mb": 0,
+        "used_percent": 0,
+    }
+
+    if TORCH_AVAILABLE and torch.cuda.is_available():
+        stats["cuda_available"] = True
+        free, total = torch.cuda.mem_get_info()
+        free_mb = free / (1024 * 1024)
+        total_mb = total / (1024 * 1024)
+        used_mb = total_mb - free_mb
+
+        stats["free_mb"] = free_mb
+        stats["total_mb"] = total_mb
+        stats["used_mb"] = used_mb
+        stats["used_percent"] = (used_mb / total_mb) * 100 if total_mb > 0 else 0
+
+    return stats
+
+
+def purge_memory(mode: str = "soft", unload_models: bool = False) -> float:
+    before_stats = get_memory_stats()
+
+    if mode == "soft":
+        # Basic garbage collection and cache clearing
+        gc.collect()
+        if TORCH_AVAILABLE and torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    elif mode == "aggressive":
+        # Multiple passes of garbage collection with full cache clearing
+        gc.collect()
+        gc.collect()
+        if TORCH_AVAILABLE and torch.cuda.is_available():
+            torch.cuda.synchronize()
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+
+    elif mode == "models_only":
+        # Only unload models
+        if COMFY_AVAILABLE:
+            mm.unload_all_models()
+            mm.soft_empty_cache()
+        gc.collect()
+
+    elif mode == "cache_only":
+        # Only clear cache without garbage collection
+        if TORCH_AVAILABLE and torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    # Handle model unloading for non-model-specific modes
+    if unload_models and mode not in ["models_only"]:
+        if COMFY_AVAILABLE:
+            mm.unload_all_models()
+            mm.soft_empty_cache()
+
+    after_stats = get_memory_stats()
+    freed_mb = before_stats["used_mb"] - after_stats["used_mb"]
+
+    return max(0, freed_mb)
+
+
+def format_memory_report(
+    before: Dict[str, float], after: Dict[str, float], mode: str, elapsed_ms: float
+) -> str:
+    if not before.get("cuda_available", True):
+        return (
+            "Memory Purge Report\n"
+            "-------------------\n"
+            "CUDA not available - CPU memory management only\n"
+            f"Mode: {mode}\n"
+            f"Time: {elapsed_ms:.1f}ms"
+        )
+
+    freed_mb = before["used_mb"] - after["used_mb"]
+
+    report = [
+        "Memory Purge Report",
+        "-------------------",
+        f"Mode: {mode}",
+        f"Memory Freed: {freed_mb:.1f} MB",
+        f"Before: {before['used_mb']:.1f} MB used ({before['used_percent']:.1f}%)",
+        f"After: {after['used_mb']:.1f} MB used ({after['used_percent']:.1f}%)",
+        f"Time: {elapsed_ms:.1f}ms",
+    ]
+
+    return "\n".join(report)
+
+
+def should_purge(threshold_mb: int) -> Tuple[bool, str]:
+    if threshold_mb <= 0:
+        return True, ""
+
+    stats = get_memory_stats()
+
+    if not stats["cuda_available"]:
+        return True, "CUDA not available, proceeding with CPU memory management"
+
+    if stats["used_mb"] >= threshold_mb:
+        return (
+            True,
+            f"Memory usage ({stats['used_mb']:.1f} MB) exceeds threshold ({threshold_mb} MB)",
+        )
+    else:
+        return (
+            False,
+            f"Memory usage ({stats['used_mb']:.1f} MB) below threshold ({threshold_mb} MB)",
+        )

kikotools/tools/kiko_purge_vram/node.py

@@ -0,0 +1,102 @@
+import time
+from typing import Any, Dict, Tuple
+
+try:
+    from ...base.base_node import ComfyAssetsBaseNode as BaseNode
+    from ...base.any_type import AnyType
+except ImportError:
+    # Fallback for testing environment
+    from kikotools.base.base_node import ComfyAssetsBaseNode as BaseNode
+    from kikotools.base.any_type import AnyType
+from .logic import get_memory_stats, purge_memory, format_memory_report, should_purge
+
+any_type = AnyType("*")
+
+
+class KikoPurgeVRAM(BaseNode):
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        return {
+            "required": {
+                "anything": (any_type, {}),
+                "mode": (
+                    ["soft", "aggressive", "models_only", "cache_only"],
+                    {
+                        "default": "soft",
+                        "tooltip": "Purge mode: soft (basic), aggressive (thorough), models_only (unload models), cache_only (clear cache)",
+                    },
+                ),
+                "report_memory": (
+                    "BOOLEAN",
+                    {
+                        "default": True,
+                        "tooltip": "Generate detailed memory usage report",
+                    },
+                ),
+            },
+            "optional": {
+                "memory_threshold_mb": (
+                    "INT",
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "max": 48000,
+                        "step": 100,
+                        "tooltip": "Only purge if memory usage exceeds this threshold (0 = always purge)",
+                    },
+                ),
+            },
+        }
+
+    RETURN_TYPES = (any_type, "STRING")
+    RETURN_NAMES = ("passthrough", "memory_report")
+    FUNCTION = "purge_vram"
+    CATEGORY = "🫶 ComfyAssets/🛠️ Utils"
+    OUTPUT_NODE = True
+    DESCRIPTION = "Purge VRAM to free up GPU memory during workflow execution. Passes through any input unchanged."
+
+    def purge_vram(
+        self,
+        anything: Any,
+        mode: str,
+        report_memory: bool,
+        memory_threshold_mb: int = 0,
+    ) -> Tuple[Any, str]:
+        # Check if we should purge based on threshold
+        should_run, threshold_msg = should_purge(memory_threshold_mb)
+
+        if not should_run:
+            if report_memory:
+                return anything, f"Memory purge skipped: {threshold_msg}"
+            else:
+                return anything, ""
+
+        # Get before stats
+        before_stats = get_memory_stats() if report_memory else None
+        start_time = time.time()
+
+        # Determine if we should unload models
+        unload_models = mode in ["models_only", "aggressive"]
+
+        # Perform memory purge
+        purge_memory(mode=mode, unload_models=unload_models)
+
+        # Calculate elapsed time
+        elapsed_ms = (time.time() - start_time) * 1000
+
+        # Generate report if requested
+        if report_memory:
+            after_stats = get_memory_stats()
+            report = format_memory_report(before_stats, after_stats, mode, elapsed_ms)
+            if threshold_msg and memory_threshold_mb > 0:
+                report = f"{threshold_msg}\n\n{report}"
+        else:
+            report = ""
+
+        # Pass through the input unchanged
+        return anything, report
+
+
+NODE_CLASS_MAPPINGS = {"KikoPurgeVRAM": KikoPurgeVRAM}
+
+NODE_DISPLAY_NAME_MAPPINGS = {"KikoPurgeVRAM": "Kiko Purge VRAM"}

tests/unit/tools/test_kiko_purge_vram.py

@@ -0,0 +1,299 @@
+import sys
+from unittest.mock import patch, MagicMock
+
+import pytest
+
+# Mock comfy modules
+sys.modules["comfy"] = MagicMock()
+sys.modules["comfy.model_management"] = MagicMock()
+
+from kikotools.tools.kiko_purge_vram.logic import (
+    purge_memory,
+    get_memory_stats,
+    format_memory_report,
+)
+
+
+class TestMemoryStats:
+    @patch("torch.cuda.is_available")
+    @patch("torch.cuda.mem_get_info")
+    def test_get_memory_stats_with_cuda(self, mock_mem_info, mock_cuda_available):
+        mock_cuda_available.return_value = True
+        mock_mem_info.return_value = (4000000000, 8000000000)  # 4GB free, 8GB total
+
+        stats = get_memory_stats()
+
+        assert stats["cuda_available"] is True
+        assert stats["free_mb"] == pytest.approx(3814.7, rel=0.1)
+        assert stats["total_mb"] == pytest.approx(7629.4, rel=0.1)
+        assert stats["used_mb"] == pytest.approx(3814.7, rel=0.1)
+        assert stats["used_percent"] == pytest.approx(50.0, rel=0.1)
+
+    @patch("torch.cuda.is_available")
+    def test_get_memory_stats_without_cuda(self, mock_cuda_available):
+        mock_cuda_available.return_value = False
+
+        stats = get_memory_stats()
+
+        assert stats["cuda_available"] is False
+        assert stats["free_mb"] == 0
+        assert stats["total_mb"] == 0
+        assert stats["used_mb"] == 0
+        assert stats["used_percent"] == 0
+
+
+class TestMemoryPurge:
+    @patch("torch.cuda.is_available")
+    @patch("torch.cuda.empty_cache")
+    @patch("torch.cuda.ipc_collect")
+    @patch("gc.collect")
+    def test_purge_memory_soft_mode(
+        self, mock_gc, mock_ipc, mock_empty_cache, mock_cuda
+    ):
+        mock_cuda.return_value = True
+
+        with patch(
+            "kikotools.tools.kiko_purge_vram.logic.get_memory_stats"
+        ) as mock_stats:
+            mock_stats.side_effect = [
+                {"used_mb": 4000, "free_mb": 4000},
+                {"used_mb": 2000, "free_mb": 6000},
+            ]
+
+            freed_mb = purge_memory(mode="soft", unload_models=False)
+
+            mock_gc.assert_called_once()
+            mock_empty_cache.assert_called_once()
+            mock_ipc.assert_not_called()
+            assert freed_mb == 2000
+
+    @patch("torch.cuda.is_available")
+    @patch("torch.cuda.empty_cache")
+    @patch("torch.cuda.ipc_collect")
+    @patch("torch.cuda.synchronize")
+    @patch("gc.collect")
+    def test_purge_memory_aggressive_mode(
+        self, mock_gc, mock_sync, mock_ipc, mock_empty_cache, mock_cuda
+    ):
+        mock_cuda.return_value = True
+
+        with patch(
+            "kikotools.tools.kiko_purge_vram.logic.get_memory_stats"
+        ) as mock_stats:
+            mock_stats.side_effect = [
+                {"used_mb": 4000, "free_mb": 4000},
+                {"used_mb": 1500, "free_mb": 6500},
+            ]
+
+            freed_mb = purge_memory(mode="aggressive", unload_models=False)
+
+            assert mock_gc.call_count == 2
+            mock_empty_cache.assert_called()
+            mock_ipc.assert_called_once()
+            mock_sync.assert_called_once()
+            assert freed_mb == 2500
+
+    @patch("kikotools.tools.kiko_purge_vram.logic.COMFY_AVAILABLE", True)
+    @patch("kikotools.tools.kiko_purge_vram.logic.mm.unload_all_models")
+    @patch("kikotools.tools.kiko_purge_vram.logic.mm.soft_empty_cache")
+    @patch("torch.cuda.is_available")
+    @patch("gc.collect")
+    def test_purge_memory_models_only(
+        self, mock_gc, mock_cuda, mock_soft_empty, mock_unload
+    ):
+        mock_cuda.return_value = True
+
+        with patch(
+            "kikotools.tools.kiko_purge_vram.logic.get_memory_stats"
+        ) as mock_stats:
+            mock_stats.side_effect = [
+                {"used_mb": 6000, "free_mb": 2000},
+                {"used_mb": 1000, "free_mb": 7000},
+            ]
+
+            freed_mb = purge_memory(mode="models_only", unload_models=True)
+
+            mock_unload.assert_called_once()
+            mock_soft_empty.assert_called_once()
+            mock_gc.assert_called()
+            assert freed_mb == 5000
+
+    @patch("torch.cuda.is_available")
+    @patch("torch.cuda.empty_cache")
+    @patch("gc.collect")
+    def test_purge_memory_cache_only(self, mock_gc, mock_empty_cache, mock_cuda):
+        mock_cuda.return_value = True
+
+        with patch(
+            "kikotools.tools.kiko_purge_vram.logic.get_memory_stats"
+        ) as mock_stats:
+            mock_stats.side_effect = [
+                {"used_mb": 3000, "free_mb": 5000},
+                {"used_mb": 2500, "free_mb": 5500},
+            ]
+
+            freed_mb = purge_memory(mode="cache_only", unload_models=False)
+
+            mock_gc.assert_not_called()
+            mock_empty_cache.assert_called_once()
+            assert freed_mb == 500
+
+    @patch("torch.cuda.is_available")
+    def test_purge_memory_no_cuda(self, mock_cuda):
+        mock_cuda.return_value = False
+
+        with patch("gc.collect") as mock_gc:
+            freed_mb = purge_memory(mode="soft", unload_models=False)
+
+            mock_gc.assert_called_once()
+            assert freed_mb == 0
+
+
+class TestMemoryReport:
+    def test_format_memory_report_with_improvement(self):
+        before = {
+            "used_mb": 4000,
+            "free_mb": 4000,
+            "total_mb": 8000,
+            "used_percent": 50,
+        }
+        after = {"used_mb": 2000, "free_mb": 6000, "total_mb": 8000, "used_percent": 25}
+
+        report = format_memory_report(before, after, mode="soft", elapsed_ms=150)
+
+        assert "Memory Purge Report" in report
+        assert "Mode: soft" in report
+        assert "Memory Freed: 2000.0 MB" in report
+        assert "Before: 4000.0 MB used (50.0%)" in report
+        assert "After: 2000.0 MB used (25.0%)" in report
+        assert "Time: 150.0ms" in report
+
+    def test_format_memory_report_no_improvement(self):
+        before = {
+            "used_mb": 2000,
+            "free_mb": 6000,
+            "total_mb": 8000,
+            "used_percent": 25,
+        }
+        after = {"used_mb": 2000, "free_mb": 6000, "total_mb": 8000, "used_percent": 25}
+
+        report = format_memory_report(before, after, mode="cache_only", elapsed_ms=50)
+
+        assert "Memory Freed: 0.0 MB" in report
+        assert "Time: 50.0ms" in report
+
+    def test_format_memory_report_no_cuda(self):
+        before = {
+            "used_mb": 0,
+            "free_mb": 0,
+            "total_mb": 0,
+            "used_percent": 0,
+            "cuda_available": False,
+        }
+        after = {
+            "used_mb": 0,
+            "free_mb": 0,
+            "total_mb": 0,
+            "used_percent": 0,
+            "cuda_available": False,
+        }
+
+        report = format_memory_report(before, after, mode="soft", elapsed_ms=10)
+
+        assert "CUDA not available" in report
+
+
+class TestKikoPurgeVRAMNode:
+    @patch("kikotools.tools.kiko_purge_vram.node.format_memory_report")
+    @patch("kikotools.tools.kiko_purge_vram.node.purge_memory")
+    @patch("kikotools.tools.kiko_purge_vram.node.get_memory_stats")
+    @patch("kikotools.tools.kiko_purge_vram.node.should_purge")
+    def test_node_execute_with_threshold(
+        self, mock_should_purge, mock_stats, mock_purge, mock_format
+    ):
+        from kikotools.tools.kiko_purge_vram.node import KikoPurgeVRAM
+
+        mock_should_purge.return_value = (
+            True,
+            "Memory usage (5000.0 MB) exceeds threshold (4000 MB)",
+        )
+        mock_stats.side_effect = [
+            {
+                "used_mb": 5000,
+                "free_mb": 3000,
+                "total_mb": 8000,
+                "used_percent": 62.5,
+                "cuda_available": True,
+            },
+            {
+                "used_mb": 2000,
+                "free_mb": 6000,
+                "total_mb": 8000,
+                "used_percent": 25,
+                "cuda_available": True,
+            },
+        ]
+        mock_purge.return_value = 3000
+        mock_format.return_value = "Memory Purge Report\n-------------------\nMode: soft\nMemory Freed: 3000.0 MB"
+
+        node = KikoPurgeVRAM()
+        test_input = "test_data"
+
+        result, report = node.purge_vram(
+            anything=test_input,
+            mode="soft",
+            report_memory=True,
+            memory_threshold_mb=4000,
+        )
+
+        assert result == test_input
+        assert "Memory Freed: 3000.0 MB" in report
+        mock_purge.assert_called_once_with(mode="soft", unload_models=False)
+
+    @patch("kikotools.tools.kiko_purge_vram.logic.get_memory_stats")
+    def test_node_skip_below_threshold(self, mock_stats):
+        from kikotools.tools.kiko_purge_vram.node import KikoPurgeVRAM
+
+        mock_stats.return_value = {
+            "used_mb": 2000,
+            "free_mb": 6000,
+            "total_mb": 8000,
+            "used_percent": 25,
+            "cuda_available": True,
+        }
+
+        node = KikoPurgeVRAM()
+        test_input = "test_data"
+
+        with patch("kikotools.tools.kiko_purge_vram.logic.purge_memory") as mock_purge:
+            result, report = node.purge_vram(
+                anything=test_input,
+                mode="soft",
+                report_memory=True,
+                memory_threshold_mb=3000,
+            )
+
+            assert result == test_input
+            assert "below threshold" in report.lower()
+            mock_purge.assert_not_called()
+
+    def test_node_input_types(self):
+        from kikotools.tools.kiko_purge_vram.node import KikoPurgeVRAM
+
+        input_types = KikoPurgeVRAM.INPUT_TYPES()
+
+        assert "required" in input_types
+        assert "optional" in input_types
+        assert "anything" in input_types["required"]
+        assert "mode" in input_types["required"]
+        assert "report_memory" in input_types["required"]
+        assert "memory_threshold_mb" in input_types["optional"]
+
+    def test_node_properties(self):
+        from kikotools.tools.kiko_purge_vram.node import KikoPurgeVRAM
+
+        assert KikoPurgeVRAM.FUNCTION == "purge_vram"
+        assert KikoPurgeVRAM.CATEGORY == "🫶 ComfyAssets/🛠️ Utils"
+        assert KikoPurgeVRAM.OUTPUT_NODE is True
+        assert len(KikoPurgeVRAM.RETURN_TYPES) == 2
+        assert KikoPurgeVRAM.RETURN_NAMES == ("passthrough", "memory_report")