Files

455 lines
17 KiB
JavaScript

import {api} from "../../../scripts/api.js";
import {app} from "../../../scripts/app.js";
import {$el} from "../../scripts/ui.js";
// region: Refresh Timer
let refreshTimer = null;
function stopRefreshTimer() {
if (!refreshTimer) {
return;
}
clearInterval(refreshTimer);
refreshTimer = null;
}
function startRefreshTimer() {
stopRefreshTimer();
refreshTimer = setInterval(function () {
app.graph.setDirtyCanvas(true, false);
}, 100);
}
// endregion
function formatExecutionTime(time) {
return `${(time / 1000.0).toFixed(2)}s`
}
// Reference: https://gist.github.com/zentala/1e6f72438796d74531803cc3833c039c
function formatBytes(bytes, decimals) {
if (bytes === 0) {
return '0 B'
}
const k = 1024,
dm = decimals || 2,
sizes = ['Bytes', 'KB', 'MB', 'GB', 'TB', 'PB', 'EB', 'ZB', 'YB'],
i = Math.floor(Math.log(bytes) / Math.log(k));
return parseFloat((bytes / Math.pow(k, i)).toFixed(dm)) + ' ' + sizes[i];
}
// Reference: https://github.com/ltdrdata/ComfyUI-Manager/blob/main/js/comfyui-manager.js
function drawBadge(node, orig, restArgs) {
let ctx = restArgs[0];
const r = orig?.apply?.(node, restArgs);
if (!node.flags.collapsed && node.constructor.title_mode != LiteGraph.NO_TITLE) {
let text = "";
if (node.ty_et_execution_time !== undefined) {
text = formatExecutionTime(node.ty_et_execution_time) + " - vram " + formatBytes(node.ty_et_vram_used, 2);
} else if (node.ty_et_start_time !== undefined) {
text = formatExecutionTime(LiteGraph.getTime() - node.ty_et_start_time);
}
if (!text) {
return
}
let fgColor = "white";
let bgColor = "#0F1F0F";
let visible = true;
ctx.save();
ctx.font = "12px sans-serif";
const textSize = ctx.measureText(text);
ctx.fillStyle = bgColor;
ctx.beginPath();
const paddingHorizontal = 6;
ctx.roundRect(0, -LiteGraph.NODE_TITLE_HEIGHT - 20, textSize.width + paddingHorizontal * 2, 20, 5);
ctx.fill();
ctx.fillStyle = fgColor;
ctx.fillText(text, paddingHorizontal, -LiteGraph.NODE_TITLE_HEIGHT - paddingHorizontal);
ctx.restore();
}
return r;
}
// Reference: https://github.com/ltdrdata/ComfyUI-Manager/blob/main/js/common.js
async function unloadModelsAndFreeMemory() {
let res = await api.fetchApi(`/free`, {
method: 'POST',
headers: {'Content-Type': 'application/json'},
body: '{"unload_models": true, "free_memory": true}'
});
if (res.status === 200) {
app.ui.dialog.show('<span style="color: green;">Unload models and free memory success.</span>')
} else {
app.ui.dialog.show('[ERROR] Unload models and free memory fail.')
}
app.ui.dialog.element.style.zIndex = 10010;
}
function setupClearExecutionCacheMenu() {
const menu = document.querySelector(".comfy-menu");
const freeButton = document.createElement("button");
freeButton.textContent = "Clear Execution Cache";
freeButton.onclick = async () => {
await unloadModelsAndFreeMemory();
};
menu.append(freeButton);
}
let lastRunningDate = null;
let runningData = null;
// https://stackoverflow.com/a/56370447
function exportTable(table, separator = ',') {
// Select rows from table_id
var rows = table.querySelectorAll('tr');
// Construct csv
var csv = [];
for (var i = 0; i < rows.length; i++) {
var row = [], cols = rows[i].querySelectorAll('td, th');
for (var j = 0; j < cols.length; j++) {
// Clean innertext to remove multiple spaces and jumpline (break csv)
var data = cols[j].innerText.replace(/(\r\n|\n|\r)/gm, '').replace(/(\s\s)/gm, ' ')
// Escape double-quote with double-double-quote (see https://stackoverflow.com/questions/17808511/properly-escape-a-double-quote-in-csv)
data = data.replace(/"/g, '""');
// Push escaped string
row.push('"' + data + '"');
}
csv.push(row.join(separator));
}
var csv_string = csv.join('\n');
// Download it
var filename = 'execution_time' + new Date().toLocaleDateString() + '.csv';
var link = document.createElement('a');
link.style.display = 'none';
link.setAttribute('target', '_blank');
link.setAttribute('href', 'data:text/csv;charset=utf-8,' + encodeURIComponent(csv_string));
link.setAttribute('download', filename);
document.body.appendChild(link);
link.click();
document.body.removeChild(link);
}
function buildTableHtml() {
const tableBody = $el("tbody")
const tableFooter = $el("tfoot", {style: {"background": "var(--comfy-input-bg)"}})
const headerThStyle = {"white-space": "nowrap"}
const table = $el("table", {
textAlign: "right",
border: "1px solid var(--border-color)",
style: {"border": "none", "border-spacing": "0", "font-size": "14px", "width": "100%"}
}, [
$el("thead", {style: {"background": "var(--comfy-input-bg)"}}, [
$el("tr", [
$el("th", {style: headerThStyle, "textContent": "Node Id"}),
$el("th", {style: headerThStyle, "textContent": "Node Title"}),
$el("th", {style: headerThStyle, "textContent": "Current Time"}),
$el("th", {style: headerThStyle, "textContent": "Per Time"}),
$el("th", {style: headerThStyle, "textContent": "Cur / Pre Time Diff"}),
$el("th", {style: headerThStyle, "textContent": "VRAM Used"})
])
]),
tableBody,
tableFooter
]);
if (runningData?.nodes_execution_time === undefined) {
return table;
}
function diff(current, pre) {
let diffText;
let diffColor;
if (pre) {
const diffTime = current - pre;
const diffPercentText = `${(diffTime * 100 / pre).toFixed(2)}%`;
if (diffTime > 0) {
diffColor = 'red';
diffText = `+${formatExecutionTime(diffTime)} / +${diffPercentText}`;
} else if (diffPercentText === '0.00%') {
diffColor = 'white';
diffText = formatExecutionTime(diffTime);
} else {
diffColor = 'green';
diffText = `${formatExecutionTime(diffTime)} / ${diffPercentText}`;
}
}
return [diffColor, diffText]
}
let max_execution_time = null
let max_vram_used = null
runningData.nodes_execution_time.forEach(function (item) {
const nodeId = item.node;
const node = app.graph.getNodeById(nodeId)
const title = node?.title ?? nodeId
const preExecutionTime = lastRunningDate?.nodes_execution_time?.find(x => x.node === nodeId)?.execution_time
const [diffColor, diffText] = diff(item.execution_time, preExecutionTime);
if (max_execution_time == null || item.execution_time > max_execution_time) {
max_execution_time = item.execution_time
}
if (max_vram_used == null || item.vram_used > max_vram_used) {
max_vram_used = item.vram_used
}
tableBody.append($el("tr", {
style: {"cursor": "pointer"},
onclick: () => {
if (node) {
app.canvas.selectNode(node, false);
}
}
}, [
$el("td", {style: {"textAlign": "right"}, "textContent": nodeId}),
$el("td", {style: {"textAlign": "right"}, "textContent": title}),
$el("td", {style: {"textAlign": "right"}, "textContent": formatExecutionTime(item.execution_time)}),
$el("td", {
style: {"textAlign": "right"},
"textContent": preExecutionTime !== undefined ? formatExecutionTime(preExecutionTime) : undefined
}),
$el("td", {
style: {
"textAlign": "right",
"color": diffColor
},
"textContent": diffText
}),
$el("td", {style: {"textAlign": "right"}, "textContent": formatBytes(item.vram_used, 2)}),
]))
});
if (runningData.total_execution_time !== null) {
const [diffColor, diffText] = diff(runningData.total_execution_time, lastRunningDate?.total_execution_time);
tableFooter.append($el("tr", [
$el("td", {style: {"textAlign": "right"}, "textContent": 'Max'}),
$el("td", {style: {"textAlign": "right"}, "textContent": ''}),
$el("td", {
style: {"textAlign": "right"},
"textContent": max_execution_time != null ? formatExecutionTime(max_execution_time) : ''
}),
$el("td", {
style: {"textAlign": "right"},
"textContent": ''
}),
$el("td", {
style: {
"textAlign": "right",
"color": diffColor
},
"textContent": ''
}),
$el("td", {
style: {"textAlign": "right"},
"textContent": max_vram_used != null ? formatBytes(max_vram_used, 2) : ''
}),
]))
tableFooter.append($el("tr", [
$el("td", {style: {"textAlign": "right"}, "textContent": 'Total'}),
$el("td", {style: {"textAlign": "right"}, "textContent": ''}),
$el("td", {
style: {"textAlign": "right"},
"textContent": formatExecutionTime(runningData.total_execution_time)
}),
$el("td", {
style: {"textAlign": "right"},
"textContent": lastRunningDate?.total_execution_time ? formatExecutionTime(lastRunningDate?.total_execution_time) : undefined
}),
$el("td", {
style: {
"textAlign": "right",
"color": diffColor
},
"textContent": diffText
}),
$el("td", {style: {"textAlign": "right"}, "textContent": ""}),
]))
}
return table;
}
function refreshTable() {
app.graph._nodes.forEach(function (node) {
if (node.comfyClass === "TY_ExecutionTime" && node.widgets) {
const tableWidget = node.widgets.find((w) => w.name === "Table");
if (!tableWidget) {
return;
}
tableWidget.inputEl.replaceChild(buildTableHtml(), tableWidget.inputEl.firstChild);
const computeSize = node.computeSize();
const newSize = [Math.max(node.size[0], computeSize[0]), Math.max(node.size[1], computeSize[1])];
node.setSize(newSize);
app.graph.setDirtyCanvas(true);
}
});
}
app.registerExtension({
name: "TyDev-Utils.ExecutionTime",
async setup() {
setupClearExecutionCacheMenu();
api.addEventListener("executing", ({detail}) => {
const nodeId = detail;
if (!nodeId) { // Finish
return
}
const node = app.graph.getNodeById(nodeId)
if (node) {
node.ty_et_start_time = LiteGraph.getTime();
}
});
api.addEventListener("TyDev-Utils.ExecutionTime.executed", ({detail}) => {
const node = app.graph.getNodeById(detail.node)
if (node) {
node.ty_et_execution_time = detail.execution_time;
node.ty_et_vram_used = detail.vram_used;
}
const index = runningData.nodes_execution_time.findIndex(x => x.node === detail.node);
const data = {
node: detail.node,
execution_time: detail.execution_time,
vram_used: detail.vram_used
};
if (index > 0) {
runningData.nodes_execution_time[index] = data
} else {
runningData.nodes_execution_time.push(data)
}
refreshTable();
});
api.addEventListener("execution_start", ({detail}) => {
if (runningData && runningData.total_execution_time == null) {
return;
}
lastRunningDate = runningData;
app.graph._nodes.forEach(function (node) {
delete node.ty_et_start_time
delete node.ty_et_execution_time
delete node.ty_et_vram_used
});
runningData = {
nodes_execution_time: [],
total_execution_time: null
};
startRefreshTimer();
});
api.addEventListener("TyDev-Utils.ExecutionTime.execution_end", ({detail}) => {
stopRefreshTimer();
runningData.total_execution_time = detail.execution_time;
refreshTable();
})
},
async nodeCreated(node, app) {
if (!node.ty_et_swizzled) {
let orig = node.onDrawForeground;
if (!orig) {
orig = node.__proto__.onDrawForeground;
}
node.onDrawForeground = function (ctx) {
drawBadge(node, orig, arguments)
};
node.ty_et_swizzled = true;
}
},
async loadedGraphNode(node, app) {
if (!node.ty_et_swizzled) {
const orig = node.onDrawForeground;
node.onDrawForeground = function (ctx) {
drawBadge(node, orig, arguments)
};
node.ty_et_swizzled = true;
}
},
async beforeRegisterNodeDef(nodeType, nodeData, app) {
if (nodeType.comfyClass === "TY_ExecutionTime") {
const originComputeSize = nodeType.prototype.computeSize || LGraphNode.prototype.computeSize;
nodeType.prototype.computeSize = function () {
const originSize = originComputeSize.apply(this, arguments);
if (this.flags?.collapsed || !this.widgets) {
return originSize;
}
const tableWidget = this.widgets.find((w) => w.name === "Table");
if (!tableWidget) {
return originSize;
}
const tableElem = tableWidget.inputEl.firstChild;
const tableHeight = tableElem.getBoundingClientRect().height;
const thHeight = tableElem.tHead.getBoundingClientRect().height;
const thUnscaledHeight = 24;
const tableUnscaledHeight = thUnscaledHeight * tableHeight / thHeight;
const autoResizeMaxHeight = 300;
return [Math.max(originSize[0], 480), originSize[1] + Math.min(tableUnscaledHeight, autoResizeMaxHeight) - LiteGraph.NODE_WIDGET_HEIGHT];
}
const nodeCreated = nodeType.prototype.onNodeCreated;
nodeType.prototype.onNodeCreated = function () {
nodeCreated?.apply(this, arguments);
const tableWidget = {
type: "HTML",
name: "Table",
draw: function (ctx, node, widgetWidth, y, widgetHeight) {
const marginHorizontal = 14;
const marginTop = 0;
const marginBottom = 14;
const elRect = ctx.canvas.getBoundingClientRect();
const transform = new DOMMatrix()
.scaleSelf(elRect.width / ctx.canvas.width, elRect.height / ctx.canvas.height)
.multiplySelf(ctx.getTransform())
.translateSelf(marginHorizontal, marginTop + y);
const x = Math.max(0, Math.round(ctx.getTransform().a * (node.size[0] - this.inputEl.scrollWidth - 2 * marginHorizontal) / 2));
Object.assign(
this.inputEl.style,
{
transformOrigin: '0 0',
transform: transform,
left: `${x}px`,
top: `0px`,
position: "absolute",
width: `${widgetWidth - marginHorizontal * 2}px`,
height: `${node.size[1] - (marginTop + marginBottom) - y}px`,
overflow: `auto`,
}
);
},
};
tableWidget.inputEl = $el("div");
document.body.appendChild(tableWidget.inputEl);
this.addWidget("button", "Export CSV", "display: none", () => {
exportTable(tableWidget.inputEl.firstChild)
});
this.addCustomWidget(tableWidget);
this.onRemoved = function () {
tableWidget.inputEl.remove();
};
this.serialize_widgets = false;
this.isVirtualNode = true;
const tableElem = buildTableHtml();
tableWidget.inputEl.appendChild(tableElem)
this.setSize(this.computeSize());
}
}
}
});