darshd9941-comfyui-memory-manager/nodes.py

"""ComfyUI custom nodes for memory management."""
import gc
import time
from .memory_manager import get_vram_info, force_cleanup, get_cache, set_cache_size


class VRAMMonitor:
    """Display real-time VRAM usage on the workflow canvas."""

    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {}, "optional": {"trigger": ("*", {})}}

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("vram_info",)
    FUNCTION = "monitor"
    CATEGORY = "memory/monitor"

    def monitor(self, trigger=None):
        info = get_vram_info()
        if not info["available"]:
            return ("No CUDA GPU available",)

        text = (
            f"VRAM: {info['allocated_mb']}MB / {info['total_mb']}MB "
            f"({info['utilization']}%)\n"
            f"Reserved: {info['reserved_mb']}MB | Free: {info['free_mb']}MB"
        )
        return (text,)


class ForceCleanup:
    """Force garbage collection and CUDA cache clear."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {},
            "optional": {
                "trigger": ("*", {}),
                "aggressive": ("BOOLEAN", {"default": False}),
            },
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("status",)
    FUNCTION = "cleanup"
    CATEGORY = "memory/management"

    def cleanup(self, trigger=None, aggressive=False):
        before = get_vram_info()
        force_cleanup()
        if aggressive:
            for _ in range(3):
                gc.collect()
            force_cleanup()
        after = get_vram_info()

        freed = before.get("allocated_mb", 0) - after.get("allocated_mb", 0)
        status = f"Freed {freed}MB VRAM" if freed > 0 else "No memory freed"
        return (status,)


class MemoryBudgetGate:
    """Gate that checks if enough VRAM is available before proceeding."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "required_mb": ("INT", {"default": 2000, "min": 100, "max": 24000, "step": 100}),
            },
            "optional": {
                "model": ("*", {}),
            },
        }

    RETURN_TYPES = ("*", "*")
    RETURN_NAMES = ("model", "status")
    FUNCTION = "check"
    CATEGORY = "memory/management"

    def check(self, required_mb=2000, model=None):
        info = get_vram_info()
        if not info["available"]:
            return (model, "No GPU — skipping check")

        if info["free_mb"] >= required_mb:
            return (model, f"OK — {info['free_mb']}MB available (need {required_mb}MB)")
        else:
            force_cleanup()
            info = get_vram_info()
            if info["free_mb"] >= required_mb:
                return (model, f"Cleaned — {info['free_mb']}MB now available")
            return (model, f"BLOCKED — only {info['free_mb']}MB free, need {required_mb}MB")


class SetCacheSize:
    """Set the maximum VRAM budget for the model cache."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "max_gb": ("FLOAT", {"default": 6.0, "min": 1.0, "max": 24.0, "step": 0.5}),
            },
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("status",)
    FUNCTION = "set_size"
    CATEGORY = "memory/management"

    def set_size(self, max_gb=6.0):
        set_cache_size(max_gb)
        return (f"Cache size set to {max_gb}GB",)


class CacheStats:
    """Show current model cache statistics."""

    @classmethod
    def INPUT_TYPES(cls):
        return {"required": {}}

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("cache_stats",)
    FUNCTION = "stats"
    CATEGORY = "memory/monitor"

    def stats(self):
        cache = get_cache()
        entries = len(cache.cache)
        size_gb = cache.usage_gb
        max_gb = cache.max_size_bytes / 1024**3

        lines = [f"Cached models: {entries}", f"Cache usage: {size_gb:.2f}GB / {max_gb:.1f}GB"]
        for key, entry in cache.cache.items():
            mb = entry["size"] / 1024**2
            lines.append(f"  - {key}: {mb:.0f}MB")

        return ("\n".join(lines),)


NODE_CLASS_MAPPINGS = {
    "VRAMMonitor": VRAMMonitor,
    "ForceCleanup": ForceCleanup,
    "MemoryBudgetGate": MemoryBudgetGate,
    "SetCacheSize": SetCacheSize,
    "CacheStats": CacheStats,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "VRAMMonitor": "📊 VRAM Monitor",
    "ForceCleanup": "🧹 Force Cleanup",
    "MemoryBudgetGate": "🚦 Memory Budget Gate",
    "SetCacheSize": "⚙️ Set Cache Size",
    "CacheStats": "📈 Cache Stats",
}