darshd9941
345f86f5ba
- VRAM Monitor node: real-time GPU memory stats - Force Cleanup node: gc + CUDA cache clear - Memory Budget Gate: block if insufficient VRAM - LRU Model Cache: auto-evict least-recently-used models - Cache Stats node: inspect loaded models - Set Cache Size node: configure VRAM budget
121 lines
3.4 KiB
Python
121 lines
3.4 KiB
Python
"""ComfyUI Memory Manager — prevent OOM crashes with smart model unloading."""
|
|
import gc
|
|
import time
|
|
import logging
|
|
from collections import OrderedDict
|
|
from typing import Optional
|
|
|
|
logger = logging.getLogger("comfyui_memory_manager")
|
|
|
|
try:
|
|
import torch
|
|
TORCH_AVAILABLE = True
|
|
except ImportError:
|
|
TORCH_AVAILABLE = False
|
|
|
|
try:
|
|
import comfy.model_management
|
|
COMFY_AVAILABLE = True
|
|
except ImportError:
|
|
COMFY_AVAILABLE = False
|
|
|
|
|
|
class LRUModelCache:
|
|
"""Least Recently Used cache for GPU models."""
|
|
|
|
def __init__(self, max_size_gb: float = 6.0):
|
|
self.max_size_bytes = int(max_size_gb * 1024**3)
|
|
self.cache: OrderedDict = OrderedDict()
|
|
self.current_size = 0
|
|
|
|
def get(self, key: str):
|
|
if key in self.cache:
|
|
self.cache.move_to_end(key)
|
|
return self.cache[key]
|
|
return None
|
|
|
|
def put(self, key: str, model, size_bytes: int):
|
|
if key in self.cache:
|
|
self.cache.move_to_end(key)
|
|
return
|
|
self.cache[key] = {"model": model, "size": size_bytes}
|
|
self.current_size += size_bytes
|
|
self._evict()
|
|
|
|
def _evict(self):
|
|
while self.current_size > self.max_size_bytes and self.cache:
|
|
key, entry = self.cache.popitem(last=False)
|
|
self.current_size -= entry["size"]
|
|
logger.info(f"Evicted model: {key} ({entry['size'] / 1024**2:.0f} MB)")
|
|
del entry["model"]
|
|
gc.collect()
|
|
if TORCH_AVAILABLE and torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
def clear(self):
|
|
self.cache.clear()
|
|
self.current_size = 0
|
|
gc.collect()
|
|
if TORCH_AVAILABLE and torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
@property
|
|
def usage_bytes(self):
|
|
return self.current_size
|
|
|
|
@property
|
|
def usage_gb(self):
|
|
return self.current_size / 1024**3
|
|
|
|
|
|
def get_vram_info() -> dict:
|
|
"""Get current GPU memory information."""
|
|
if not TORCH_AVAILABLE or not torch.cuda.is_available():
|
|
return {"available": False, "total_mb": 0, "used_mb": 0, "free_mb": 0}
|
|
|
|
total = torch.cuda.get_device_properties(0).total_mem / 1024**2
|
|
reserved = torch.cuda.memory_reserved(0) / 1024**2
|
|
allocated = torch.cuda.memory_allocated(0) / 1024**2
|
|
free = total - reserved
|
|
|
|
return {
|
|
"available": True,
|
|
"total_mb": round(total),
|
|
"reserved_mb": round(reserved),
|
|
"allocated_mb": round(allocated),
|
|
"free_mb": round(free),
|
|
"utilization": round(allocated / total * 100, 1),
|
|
}
|
|
|
|
|
|
def force_cleanup():
|
|
"""Force garbage collection and CUDA cache clear."""
|
|
gc.collect()
|
|
if TORCH_AVAILABLE and torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
torch.cuda.synchronize()
|
|
|
|
|
|
def estimate_model_size(model) -> Optional[int]:
|
|
"""Estimate model VRAM usage in bytes."""
|
|
if not TORCH_AVAILABLE:
|
|
return None
|
|
try:
|
|
if hasattr(model, "model"):
|
|
return sum(p.numel() * p.element_size() for p in model.model.parameters())
|
|
return sum(p.numel() * p.element_size() for p in model.parameters())
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
# Global cache instance
|
|
_model_cache = LRUModelCache(max_size_gb=6.0)
|
|
|
|
|
|
def get_cache() -> LRUModelCache:
|
|
return _model_cache
|
|
|
|
|
|
def set_cache_size(gb: float):
|
|
_model_cache.max_size_bytes = int(gb * 1024**3)
|