diff --git a/README.md b/README.md index 0c234b9..d4fbc3d 100644 --- a/README.md +++ b/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:** diff --git a/kikotools/__init__.py b/kikotools/__init__.py index 11702a0..e18a28b 100644 --- a/kikotools/__init__.py +++ b/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", diff --git a/kikotools/base/any_type.py b/kikotools/base/any_type.py new file mode 100644 index 0000000..c476766 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/kikotools/tools/kiko_purge_vram/__init__.py b/kikotools/tools/kiko_purge_vram/__init__.py new file mode 100644 index 0000000..6a93d83 --- /dev/null +++ b/kikotools/tools/kiko_purge_vram/__init__.py @@ -0,0 +1,3 @@ +from .node import KikoPurgeVRAM + +__all__ = ["KikoPurgeVRAM"] diff --git a/kikotools/tools/kiko_purge_vram/logic.py b/kikotools/tools/kiko_purge_vram/logic.py new file mode 100644 index 0000000..942e1a9 --- /dev/null +++ b/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)", + ) diff --git a/kikotools/tools/kiko_purge_vram/node.py b/kikotools/tools/kiko_purge_vram/node.py new file mode 100644 index 0000000..b019017 --- /dev/null +++ b/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"} diff --git a/tests/unit/tools/test_kiko_purge_vram.py b/tests/unit/tools/test_kiko_purge_vram.py new file mode 100644 index 0000000..4ee0c09 --- /dev/null +++ b/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")