Real-time memory monitor for ComfyUI on Apple Silicon (MPS / unified memory): floating panel (Unified RAM, MPS driver/active, MLX Metal, Swap, CPU, process RSS, loaded models) with per-run peak markers, plus a Mac Memory Probe graph node for inline snapshots. MLX memory is read only from the execution thread for thread-safety. Includes an mflux AnyModel example workflow. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
366 lines
13 KiB
JavaScript
366 lines
13 KiB
JavaScript
import { app } from "../../scripts/app.js";
|
|
import { api } from "../../scripts/api.js";
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// ComfyUI-MacMemoryMonitor — live floating panel for Apple Silicon (unified
|
|
// memory / MPS). Polls /macmem/stats and draws DOM bars. No canvas, no themes,
|
|
// no CUDA/NVML — deliberately small.
|
|
// ---------------------------------------------------------------------------
|
|
|
|
const STORAGE_KEY = "macmem_monitor_state";
|
|
const POLL_MS = 500;
|
|
|
|
console.info("[MacMem] extension module loaded");
|
|
|
|
// ComfyUI auto-injects extension .js files but NOT .css — load our stylesheet
|
|
// ourselves, or the panel renders unstyled (no position:fixed) and is invisible.
|
|
(function loadStylesheet() {
|
|
const id = "macmem-monitor-stylesheet";
|
|
if (document.getElementById(id)) return;
|
|
const link = document.createElement("link");
|
|
link.id = id;
|
|
link.rel = "stylesheet";
|
|
link.href = new URL("./mac_mem_monitor.css", import.meta.url).href;
|
|
document.head.appendChild(link);
|
|
console.info("[MacMem] stylesheet injected:", link.href);
|
|
})();
|
|
|
|
function loadState() {
|
|
try { return JSON.parse(localStorage.getItem(STORAGE_KEY)) || {}; }
|
|
catch { return {}; }
|
|
}
|
|
function saveState(patch) {
|
|
const s = loadState();
|
|
Object.assign(s, patch);
|
|
try { localStorage.setItem(STORAGE_KEY, JSON.stringify(s)); } catch {}
|
|
}
|
|
|
|
function fmtBytes(n) {
|
|
if (n == null || isNaN(n)) return "—";
|
|
const gb = n / (1024 ** 3);
|
|
if (gb >= 1) return gb.toFixed(2) + " GB";
|
|
const mb = n / (1024 ** 2);
|
|
return mb.toFixed(0) + " MB";
|
|
}
|
|
function pct(used, total) {
|
|
if (!total || total <= 0) return 0;
|
|
return Math.max(0, Math.min(100, (used / total) * 100));
|
|
}
|
|
|
|
// --- panel construction ----------------------------------------------------
|
|
let panel, bodyEl, titleEl;
|
|
let pollTimer = null;
|
|
let running = false;
|
|
// peak-during-run tracking (client-side; reset on execution_start)
|
|
let peak = { ram: 0, mpsDriver: 0, mlx: 0 };
|
|
|
|
function makeBar(label) {
|
|
const wrap = document.createElement("div");
|
|
wrap.className = "macmem-bar-wrap";
|
|
|
|
const head = document.createElement("div");
|
|
head.className = "macmem-bar-head";
|
|
const name = document.createElement("span");
|
|
name.className = "macmem-bar-name";
|
|
name.textContent = label;
|
|
const val = document.createElement("span");
|
|
val.className = "macmem-bar-val";
|
|
head.append(name, val);
|
|
|
|
const track = document.createElement("div");
|
|
track.className = "macmem-bar-track";
|
|
const fill = document.createElement("div");
|
|
fill.className = "macmem-bar-fill";
|
|
const peakMark = document.createElement("div");
|
|
peakMark.className = "macmem-bar-peak";
|
|
track.append(fill, peakMark);
|
|
|
|
wrap.append(head, track);
|
|
wrap.__parts = { val, fill, peakMark };
|
|
return wrap;
|
|
}
|
|
|
|
function createPanel() {
|
|
const st = loadState();
|
|
|
|
panel = document.createElement("div");
|
|
panel.className = "macmem-panel";
|
|
// Default to the open canvas area (clear of the left sidebar), and clamp any
|
|
// previously-saved position back into the viewport so it can't get lost
|
|
// off-screen or hidden behind the sidebar.
|
|
const vw = window.innerWidth || 1280;
|
|
const vh = window.innerHeight || 800;
|
|
let left = st.left ?? 360;
|
|
let top = st.top ?? 90;
|
|
left = Math.min(Math.max(0, left), Math.max(0, vw - 120));
|
|
top = Math.min(Math.max(0, top), Math.max(0, vh - 60));
|
|
panel.style.left = left + "px";
|
|
panel.style.top = top + "px";
|
|
if (st.collapsed) panel.classList.add("macmem-collapsed");
|
|
// Intentionally NOT honoring a persisted "hidden" flag on load: some ComfyUI
|
|
// builds don't re-fire the settings onChange at startup, which could leave the
|
|
// panel stuck invisible with no way back. It always shows on load; the ✕ / the
|
|
// settings toggle hide it for the current session only.
|
|
|
|
// header (drag handle + collapse + close)
|
|
const header = document.createElement("div");
|
|
header.className = "macmem-header";
|
|
titleEl = document.createElement("span");
|
|
titleEl.className = "macmem-title";
|
|
titleEl.textContent = "Mac Memory";
|
|
const spacer = document.createElement("span");
|
|
spacer.style.flex = "1";
|
|
const collapseBtn = document.createElement("button");
|
|
collapseBtn.className = "macmem-btn";
|
|
collapseBtn.textContent = st.collapsed ? "▸" : "▾";
|
|
collapseBtn.title = "Collapse / expand";
|
|
const closeBtn = document.createElement("button");
|
|
closeBtn.className = "macmem-btn";
|
|
closeBtn.textContent = "✕";
|
|
closeBtn.title = "Hide (re-open from the ⚙ settings → Mac Memory Monitor)";
|
|
header.append(titleEl, spacer, collapseBtn, closeBtn);
|
|
|
|
bodyEl = document.createElement("div");
|
|
bodyEl.className = "macmem-body";
|
|
|
|
// fixed bars
|
|
bodyEl.__ram = makeBar("Unified RAM");
|
|
bodyEl.__mps = makeBar("MPS driver");
|
|
bodyEl.__active = makeBar("MPS active");
|
|
bodyEl.__mlx = makeBar("MLX Metal"); // MLX allocator (mflux etc.); hidden when MLX absent
|
|
bodyEl.__swap = makeBar("Swap");
|
|
bodyEl.append(bodyEl.__ram, bodyEl.__mps, bodyEl.__active, bodyEl.__mlx, bodyEl.__swap);
|
|
|
|
// stats line
|
|
bodyEl.__meta = document.createElement("div");
|
|
bodyEl.__meta.className = "macmem-meta";
|
|
bodyEl.append(bodyEl.__meta);
|
|
|
|
// models section
|
|
bodyEl.__models = document.createElement("div");
|
|
bodyEl.__models.className = "macmem-models";
|
|
bodyEl.append(bodyEl.__models);
|
|
|
|
panel.append(header, bodyEl);
|
|
document.body.appendChild(panel);
|
|
|
|
// --- interactions ---
|
|
collapseBtn.onclick = () => {
|
|
const collapsed = panel.classList.toggle("macmem-collapsed");
|
|
collapseBtn.textContent = collapsed ? "▸" : "▾";
|
|
saveState({ collapsed });
|
|
};
|
|
closeBtn.onclick = () => hidePanel();
|
|
|
|
// drag
|
|
let dragging = false, sx = 0, sy = 0, ox = 0, oy = 0;
|
|
header.addEventListener("pointerdown", (e) => {
|
|
if (e.target.classList.contains("macmem-btn")) return;
|
|
dragging = true;
|
|
sx = e.clientX; sy = e.clientY;
|
|
ox = panel.offsetLeft; oy = panel.offsetTop;
|
|
header.setPointerCapture(e.pointerId);
|
|
e.preventDefault();
|
|
});
|
|
header.addEventListener("pointermove", (e) => {
|
|
if (!dragging) return;
|
|
const left = Math.max(0, ox + e.clientX - sx);
|
|
const top = Math.max(0, oy + e.clientY - sy);
|
|
panel.style.left = left + "px";
|
|
panel.style.top = top + "px";
|
|
});
|
|
header.addEventListener("pointerup", (e) => {
|
|
if (!dragging) return;
|
|
dragging = false;
|
|
header.releasePointerCapture(e.pointerId);
|
|
saveState({ left: panel.offsetLeft, top: panel.offsetTop });
|
|
});
|
|
}
|
|
|
|
function hidePanel() {
|
|
if (panel) panel.style.display = "none";
|
|
saveState({ hidden: true });
|
|
stopPolling();
|
|
}
|
|
function showPanel() {
|
|
if (!panel) createPanel();
|
|
panel.style.display = "";
|
|
saveState({ hidden: false });
|
|
startPolling();
|
|
}
|
|
|
|
// --- rendering -------------------------------------------------------------
|
|
function setBar(barEl, used, total, { peakVal = null, warnPct = 90 } = {}) {
|
|
const p = pct(used, total);
|
|
const parts = barEl.__parts;
|
|
parts.fill.style.width = p.toFixed(1) + "%";
|
|
parts.fill.classList.toggle("macmem-warn", p >= warnPct);
|
|
parts.val.textContent = `${fmtBytes(used)} / ${fmtBytes(total)}`;
|
|
if (peakVal != null && total > 0) {
|
|
parts.peakMark.style.display = "block";
|
|
parts.peakMark.style.left = pct(peakVal, total).toFixed(1) + "%";
|
|
} else {
|
|
parts.peakMark.style.display = "none";
|
|
}
|
|
}
|
|
|
|
function render(d) {
|
|
if (d.error) {
|
|
bodyEl.__meta.textContent = "error: " + d.error;
|
|
return;
|
|
}
|
|
const mps = d.mps || {};
|
|
const mlx = d.mlx || null;
|
|
const mlxTotal = mlx ? ((mlx.active || 0) + (mlx.cache || 0)) : 0;
|
|
|
|
// track peaks during a run
|
|
if (running) {
|
|
peak.ram = Math.max(peak.ram, d.used_ram || 0);
|
|
peak.mpsDriver = Math.max(peak.mpsDriver, mps.driver_allocated || 0);
|
|
peak.mlx = Math.max(peak.mlx || 0, mlxTotal);
|
|
}
|
|
|
|
setBar(bodyEl.__ram, d.used_ram, d.total_ram, { peakVal: running ? peak.ram : null });
|
|
|
|
// MPS driver-allocated vs the recommended-max "budget" watermark
|
|
const budget = mps.recommended_max || d.total_ram;
|
|
setBar(bodyEl.__mps, mps.driver_allocated, budget,
|
|
{ peakVal: running ? peak.mpsDriver : null });
|
|
setBar(bodyEl.__active, mps.current_allocated, budget);
|
|
|
|
// MLX allocator (active + cache) vs the same budget — only when MLX is present.
|
|
if (mlx) {
|
|
bodyEl.__mlx.style.display = "";
|
|
setBar(bodyEl.__mlx, mlxTotal, budget, { peakVal: running ? peak.mlx : null });
|
|
} else {
|
|
bodyEl.__mlx.style.display = "none";
|
|
}
|
|
|
|
setBar(bodyEl.__swap, d.used_swap, d.total_swap, { warnPct: 50 });
|
|
|
|
// meta line
|
|
const metaBits = [];
|
|
if (d.cpu_util != null) metaBits.push(`CPU ${d.cpu_util.toFixed(0)}%`);
|
|
if (d.process_ram != null) metaBits.push(`Comfy RSS ${fmtBytes(d.process_ram)}`);
|
|
if (mlx) metaBits.push(`MLX cache ${fmtBytes(mlx.cache || 0)}`);
|
|
if (mps.recommended_max) metaBits.push(`budget ${fmtBytes(mps.recommended_max)}`);
|
|
bodyEl.__meta.textContent = metaBits.join(" · ");
|
|
|
|
// running indicator on the title
|
|
titleEl.textContent = running ? "Mac Memory ●" : "Mac Memory";
|
|
titleEl.classList.toggle("macmem-running", running);
|
|
|
|
// models
|
|
renderModels(d.models || []);
|
|
}
|
|
|
|
function renderModels(models) {
|
|
const box = bodyEl.__models;
|
|
if (!models.length) {
|
|
box.innerHTML = `<div class="macmem-models-empty">no models loaded</div>`;
|
|
return;
|
|
}
|
|
let html = `<div class="macmem-models-title">Loaded models (${models.length})</div>`;
|
|
for (const m of models) {
|
|
const size = m.loaded_size != null ? fmtBytes(m.loaded_size)
|
|
: (m.total_size != null ? fmtBytes(m.total_size) : "—");
|
|
const dev = m.device ? ` <span class="macmem-dev">${m.device}</span>` : "";
|
|
html += `<div class="macmem-model-row"><span class="macmem-model-name">${m.name}</span>`
|
|
+ `<span class="macmem-model-size">${size}${dev}</span></div>`;
|
|
}
|
|
box.innerHTML = html;
|
|
}
|
|
|
|
// --- polling ---------------------------------------------------------------
|
|
// Try several ways to reach the route so we don't depend on api.fetchApi's
|
|
// (version-dependent) base-URL handling.
|
|
async function fetchStats() {
|
|
const attempts = [];
|
|
if (api && typeof api.fetchApi === "function") {
|
|
attempts.push(() => api.fetchApi("/macmem/stats"));
|
|
}
|
|
attempts.push(() => fetch("/macmem/stats"));
|
|
attempts.push(() => fetch("./macmem/stats"));
|
|
let lastErr;
|
|
for (const go of attempts) {
|
|
try {
|
|
const r = await go();
|
|
if (r && r.ok) return await r.json();
|
|
lastErr = new Error("HTTP " + (r ? r.status : "?"));
|
|
} catch (e) { lastErr = e; }
|
|
}
|
|
throw lastErr || new Error("no fetch path worked");
|
|
}
|
|
|
|
async function poll() {
|
|
let d;
|
|
try {
|
|
d = await fetchStats();
|
|
} catch (e) {
|
|
console.error("[MacMem] poll failed:", e);
|
|
if (bodyEl) bodyEl.__meta.textContent = "poll failed: " + (e && e.message || e);
|
|
return;
|
|
}
|
|
try {
|
|
render(d);
|
|
} catch (e) {
|
|
console.error("[MacMem] render failed:", e, "data:", d);
|
|
if (bodyEl) bodyEl.__meta.textContent = "render error: " + (e && e.message || e);
|
|
}
|
|
}
|
|
function startPolling() {
|
|
if (pollTimer) return;
|
|
console.info("[MacMem] polling started");
|
|
poll();
|
|
pollTimer = setInterval(poll, POLL_MS);
|
|
}
|
|
function stopPolling() {
|
|
if (pollTimer) { clearInterval(pollTimer); pollTimer = null; }
|
|
}
|
|
|
|
// --- run state (drives peak reset + live indicator) ------------------------
|
|
function hookExecEvents() {
|
|
api.addEventListener("execution_start", () => {
|
|
running = true;
|
|
peak = { ram: 0, mpsDriver: 0, mlx: 0 };
|
|
});
|
|
api.addEventListener("execution_success", () => { running = false; });
|
|
api.addEventListener("execution_error", () => { running = false; });
|
|
api.addEventListener("execution_interrupted", () => { running = false; });
|
|
// status carries queue_remaining; when it drains, we're idle
|
|
api.addEventListener("status", (e) => {
|
|
const rem = e?.detail?.exec_info?.queue_remaining;
|
|
if (rem === 0) running = false;
|
|
});
|
|
}
|
|
|
|
// --- register --------------------------------------------------------------
|
|
app.registerExtension({
|
|
name: "MacMemoryMonitor",
|
|
async setup() {
|
|
console.info("[MacMem] setup() running");
|
|
try {
|
|
createPanel();
|
|
} catch (e) {
|
|
console.error("[MacMem] createPanel failed:", e);
|
|
return;
|
|
}
|
|
hookExecEvents();
|
|
startPolling(); // always poll on load; visibility is session-only
|
|
|
|
// Settings toggle to re-open the panel after closing it. Wrapped so a
|
|
// settings-API change in a future frontend can't break the monitor.
|
|
try {
|
|
app.ui.settings.addSetting({
|
|
id: "MacMemoryMonitor.show",
|
|
name: "Mac Memory Monitor: show panel",
|
|
type: "boolean",
|
|
defaultValue: true,
|
|
onChange: (v) => { v ? showPanel() : hidePanel(); },
|
|
});
|
|
} catch (e) {
|
|
console.warn("[MacMem] addSetting failed (non-fatal):", e);
|
|
}
|
|
},
|
|
});
|