Dreamspire's picture
custom_nodes
f2dbf59
import { api } from "../../scripts/api.js";
// copying from https://github.com/pythongosssss/ComfyUI-WD14-Tagger
class InspireProgressBadge {
constructor() {
if (!window.__progress_badge__) {
window.__progress_badge__ = Symbol("__inspire_progress_badge__");
}
this.symbol = window.__progress_badge__;
}
getState(node) {
return node[this.symbol] || {};
}
setState(node, state) {
node[this.symbol] = state;
app.canvas.setDirty(true);
}
addStatusHandler(nodeType) {
if (nodeType[this.symbol]?.statusTagHandler) {
return;
}
if (!nodeType[this.symbol]) {
nodeType[this.symbol] = {};
}
nodeType[this.symbol] = {
statusTagHandler: true,
};
api.addEventListener("inspire/update_status", ({ detail }) => {
let { node, progress, text } = detail;
const n = app.graph.getNodeById(+(node || app.runningNodeId));
if (!n) return;
const state = this.getState(n);
state.status = Object.assign(state.status || {}, { progress: text ? progress : null, text: text || null });
this.setState(n, state);
});
const self = this;
const onDrawForeground = nodeType.prototype.onDrawForeground;
nodeType.prototype.onDrawForeground = function (ctx) {
const r = onDrawForeground?.apply?.(this, arguments);
const state = self.getState(this);
if (!state?.status?.text) {
return r;
}
const { fgColor, bgColor, text, progress, progressColor } = { ...state.status };
ctx.save();
ctx.font = "12px sans-serif";
const sz = ctx.measureText(text);
ctx.fillStyle = bgColor || "dodgerblue";
ctx.beginPath();
ctx.roundRect(0, -LiteGraph.NODE_TITLE_HEIGHT - 20, sz.width + 12, 20, 5);
ctx.fill();
if (progress) {
ctx.fillStyle = progressColor || "green";
ctx.beginPath();
ctx.roundRect(0, -LiteGraph.NODE_TITLE_HEIGHT - 20, (sz.width + 12) * progress, 20, 5);
ctx.fill();
}
ctx.fillStyle = fgColor || "#fff";
ctx.fillText(text, 6, -LiteGraph.NODE_TITLE_HEIGHT - 6);
ctx.restore();
return r;
};
}
}
export const inspireProgressBadge = new InspireProgressBadge();