stratopause-lscandClaude Opus 4.8
b7c3cc181d
The /macmem/stats route previously never read MLX and returned a cached value populated only by the MacMemProbe node, so the "MLX Metal" bar stayed hidden (display:none) until a probe ran in the graph. Read MLX memory live in collect_stats() whenever the server is idle (mlx_live or not _is_executing()), keeping the last good value during execution. Reading MLX's memory counters (get_active/cache/peak_memory) is a global per-process query that is safe from any thread — verified under concurrent GPU compute — and gating the route read on "not executing" ensures MLX is never touched while a prompt runs on the executor thread. The Metal bar now appears automatically after any MLX/mflux activity. MacMemProbe still works and is now optional (live updates during a run). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
349 lines
12 KiB
Python
349 lines
12 KiB
Python
"""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"]
|