"""ComfyUI-MacMemoryMonitor Real-time memory monitor for ComfyUI on Apple Silicon (MPS / unified memory). Unlike Nvidia setups, a Mac has no separate VRAM pool: the GPU and CPU share one unified memory bank. So instead of "VRAM vs RAM", this reports: - Unified memory (system RAM) total / used -- from psutil - MPS driver-allocated memory -- torch.mps.driver_allocated_memory() - MPS currently-allocated (active) memory -- torch.mps.current_allocated_memory() - MPS recommended max ("budget" watermark) -- torch.mps.recommended_max_memory() - Swap total / used -- psutil - This process' resident set size (RSS) -- psutil - ComfyUI's currently loaded models + sizes -- comfy.model_management It registers: - an aiohttp GET route /macmem/stats returning the JSON above - a web/ JS extension that draws a live floating panel - a graph node "Mac Memory Probe" (MacMemProbe) — a pass-through you can wire inline anywhere to log a memory snapshot at that point in the run. """ import logging import platform import subprocess import psutil import torch import server from aiohttp import web import comfy.model_management as mm log = logging.getLogger(__name__) WEB_DIRECTORY = "web" # --------------------------------------------------------------------------- # # Static hardware info (cached — doesn't change at runtime) # --------------------------------------------------------------------------- # _cpu_name_cache = {"tried": False, "name": None} def _get_cpu_name(): if _cpu_name_cache["tried"]: return _cpu_name_cache["name"] _cpu_name_cache["tried"] = True try: if platform.system() == "Darwin": _cpu_name_cache["name"] = subprocess.check_output( ["sysctl", "-n", "machdep.cpu.brand_string"], timeout=2 ).decode().strip() else: _cpu_name_cache["name"] = platform.processor() or None except Exception as e: # noqa: BLE001 log.debug("macmem: cpu name lookup failed: %s", e) return _cpu_name_cache["name"] # --------------------------------------------------------------------------- # # MPS memory # --------------------------------------------------------------------------- # def _mps_stats(): """Return MPS memory figures, or None when MPS isn't the active backend.""" try: if not torch.backends.mps.is_available(): return None except Exception: # noqa: BLE001 return None def _safe(fn): try: return int(fn()) except Exception: # noqa: BLE001 return None return { "driver_allocated": _safe(torch.mps.driver_allocated_memory), "current_allocated": _safe(torch.mps.current_allocated_memory), "recommended_max": _safe(torch.mps.recommended_max_memory), } def _mlx_stats(): """MLX's own Metal allocator (used by mflux and other MLX nodes). This is invisible to torch.mps.current_allocated_memory() because MLX allocates Metal buffers directly, bypassing PyTorch's allocator. MLX's counters are PER-PROCESS, so this only reports real numbers when called from inside the ComfyUI server process (which is where this route runs). - active: bytes held by live MLX arrays right now - cache: bytes freed but retained by MLX for reuse (this is what grows across repeated mflux runs and can march toward the budget) - peak: high-water mark since process start / last reset """ try: import mlx.core as mx except Exception: return None def _safe(name): fn = getattr(mx, name, None) if fn is None: return None try: return int(fn()) except Exception: # noqa: BLE001 return None active = _safe("get_active_memory") cache = _safe("get_cache_memory") peak = _safe("get_peak_memory") if active is None and cache is None and peak is None: return None return {"active": active, "cache": cache, "peak": peak} # --------------------------------------------------------------------------- # # Loaded models (mirrors comfy's model_management internals, defensively) # --------------------------------------------------------------------------- # def _loaded_models(): models = [] try: loaded = list(mm.current_loaded_models) except Exception as e: # noqa: BLE001 log.debug("macmem: current_loaded_models failed: %s", e) return models for idx, lm in enumerate(loaded): try: patcher = getattr(lm, "model", None) if patcher is None: continue model_obj = getattr(patcher, "model", None) name = model_obj.__class__.__name__ if model_obj is not None else "unknown" try: total_size = int(patcher.model_size()) except Exception: # noqa: BLE001 total_size = None try: loaded_size = int(patcher.loaded_size()) except Exception: # noqa: BLE001 loaded_size = None device = None for attr in ("load_device", "current_device", "device"): d = getattr(patcher, attr, None) if d is not None: device = str(d) break models.append({ "index": idx, "name": name, "total_size": total_size, "loaded_size": loaded_size, "device": device, }) except Exception as e: # noqa: BLE001 log.debug("macmem: model entry %s failed: %s", idx, e) return models def _is_executing(): try: return bool(server.PromptServer.instance.prompt_queue.currently_running) except Exception: # noqa: BLE001 return False # Last MLX snapshot, captured ONLY from the ComfyUI execution thread (the probe # node). MLX uses per-thread GPU streams, so calling mlx.core memory functions # from the aiohttp/web-server thread while a model runs on the executor thread # throws "There is no Stream(gpu, 0) in current thread" — an uncatchable C++ # abort that crashes the whole process. So the HTTP route NEVER calls MLX; it # returns this cached value instead, refreshed each time the probe node runs. _last_mlx = None def collect_stats(mlx_live=False): """Single source of truth — used by both the HTTP route and the probe node. mlx_live=True (probe node, execution thread) forces a live MLX read. The HTTP route (mlx_live=False) reads MLX live too, but only while the server is idle (never concurrently with a running prompt) — so the Metal bar appears without a probe. During execution the route returns the last cached value. """ global _last_mlx # Refresh the MLX reading when explicitly asked (probe node, execution # thread) OR whenever the server is idle. Reading MLX memory counters is a # global per-process allocator query that is safe to call from any thread # (verified under concurrent GPU compute); gating the web-route read on # "not executing" additionally guarantees we never touch MLX while a model # is running on the executor thread. Net effect: the "MLX Metal" bar becomes # visible automatically after any MLX/mflux activity, with no probe node # wired inline. Keep the last good value rather than clobbering it with None. if mlx_live or not _is_executing(): _mlx = _mlx_stats() if _mlx is not None: _last_mlx = _mlx try: device = str(mm.get_torch_device()) except Exception: # noqa: BLE001 device = None vm = psutil.virtual_memory() try: sw = psutil.swap_memory() swap_total, swap_used = sw.total, sw.used except Exception: # noqa: BLE001 swap_total = swap_used = 0 try: process_ram = psutil.Process().memory_info().rss except Exception: # noqa: BLE001 process_ram = None try: cpu_util = psutil.cpu_percent(interval=None) except Exception: # noqa: BLE001 cpu_util = None return { "device": device, "cpu_name": _get_cpu_name(), "cpu_util": cpu_util, "total_ram": vm.total, "used_ram": vm.used, "available_ram": vm.available, "total_swap": swap_total, "used_swap": swap_used, "process_ram": process_ram, "mps": _mps_stats(), "mlx": _last_mlx, # live when idle (or from probe); cached during execution "models": _loaded_models(), "executing": _is_executing(), } # --------------------------------------------------------------------------- # # HTTP route # --------------------------------------------------------------------------- # try: _routes = server.PromptServer.instance.routes @_routes.get("/macmem/stats") async def _macmem_stats(request): # noqa: ARG001 try: return web.json_response(collect_stats()) except Exception as e: # noqa: BLE001 log.exception("macmem: stats collection failed") return web.json_response({"error": str(e)}, status=500) log.info("macmem: registered /macmem/stats route") except Exception as e: # noqa: BLE001 log.warning("macmem: could not register route: %s", e) # --------------------------------------------------------------------------- # # Graph node: inline memory probe (pass-through) # --------------------------------------------------------------------------- # class _AnyType(str): """Type that compares equal to any other type, so this node's socket accepts (and returns) any wire — image, latent, model, conditioning, etc.""" def __ne__(self, other): # noqa: D401 return False _ANY = _AnyType("*") def _fmt_gb(n): if n is None: return "n/a" return f"{n / (1024 ** 3):.2f} GB" class MacMemProbe: """Pass-through node that logs a memory snapshot when it executes. Drop it inline on any wire (e.g. right after the KSampler, or between VAE decode stages) to see unified-memory / MPS usage at that exact point in the run, printed to the ComfyUI console and shown on the node. """ @classmethod def INPUT_TYPES(cls): return { "required": { "input": (_ANY, {}), }, "optional": { "label": ("STRING", {"default": "probe", "multiline": False}), }, } RETURN_TYPES = (_ANY, "STRING") RETURN_NAMES = ("output", "report") FUNCTION = "probe" CATEGORY = "utils/memory" OUTPUT_NODE = True def probe(self, input, label="probe"): # Runs on the ComfyUI execution thread — safe to read MLX live here. s = collect_stats(mlx_live=True) mps = s.get("mps") or {} mlx = s.get("mlx") lines = [ f"[MacMemProbe:{label}]", f" Unified RAM : {_fmt_gb(s['used_ram'])} / {_fmt_gb(s['total_ram'])} used", f" MPS driver : {_fmt_gb(mps.get('driver_allocated'))} allocated" f" (active {_fmt_gb(mps.get('current_allocated'))}," f" budget {_fmt_gb(mps.get('recommended_max'))})", ] if mlx is not None: lines.append( f" MLX Metal : {_fmt_gb((mlx.get('active') or 0) + (mlx.get('cache') or 0))}" f" (active {_fmt_gb(mlx.get('active'))}, cache {_fmt_gb(mlx.get('cache'))}," f" peak {_fmt_gb(mlx.get('peak'))})" ) lines += [ f" Swap : {_fmt_gb(s['used_swap'])} / {_fmt_gb(s['total_swap'])}", f" Process RSS : {_fmt_gb(s['process_ram'])}", f" Models : {len(s['models'])} loaded", ] report = "\n".join(lines) log.info("\n" + report) print(report, flush=True) return { "ui": {"text": [report]}, "result": (input, report), } NODE_CLASS_MAPPINGS = { "MacMemProbe": MacMemProbe, } NODE_DISPLAY_NAME_MAPPINGS = { "MacMemProbe": "Mac Memory Probe", } __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]