IMGLAB-net-IMGLAB_ComfyUI

const MAX_MASKS = 8; const MASK_TYPE = "MASK"; const RES_COLOR = "#aaaaaa"; const OK_COLOR = "#22cc66"; const ERR_COLOR = "#cc3333"; const NA_COLOR = "#555555"; const ELSE_COLOR = "#cc9900";

const SLOT_FONT = "12px Arial"; const SLOT_NAME_X = 14; const DOT_RADIUS = 5; const NAME_GAP = 12; // gap: end of slot name → dot centre const DOT_GAP = 10; // gap: dot right edge → resolution text left edge

// LiteGraph draws output labels right-aligned from nodeW-10. // Our widest output name is "mask" or "rgba" (~30px in 12px Arial). // We leave 14px extra buffer, so output area = nodeW - 10 - 30 - 14 = nodeW - 54. // Resolution text right edge must stay at or left of this boundary. const OUTPUT_LABEL_RESERVE = 54;

const HELP_SECTIONS = [ { title: "Inputs", items: [ ["image", "Defines the target resolution for matching."], ["mask_1..8", "Connect as many masks as needed. A new slot appears as each one is connected."], ["else", "Fallback mask. Visible only when no_match is set to Use Else Mask."], ] }, { title: "Matching", items: [ ["✓ green", "Mask dimensions match the image exactly."], ["✗ red", "Mask dimensions do not match."], ["priority", "The lowest-numbered matching slot always wins."], ] }, { title: "multi_match", items: [ ["Use First Match", "Silently picks the lowest matching slot."], ["Raise Error", "Halts the queue and lists the conflicting slot numbers."], ] }, { title: "no_match", items: [ ["Output Empty Mask", "Outputs a zeroed mask (RGBA will be fully transparent)."], ["Use Else Mask", "Uses the else input as the fallback output."], ["Raise Error", "Halts the queue and lists the sizes that were connected."], ] }, { title: "Outputs", items: [ ["image", "Direct passthrough of the input image."], ["mask", "The winning mask (or fallback)."], ["rgba", "Input image + matched mask as alpha channel [B,H,W,4]."], ] }, ];

const HELP_BTN_W = 16; const HELP_BTN_H = 14; const HELP_BTN_R = 6;

let _helpPanel = null;

app.registerExtension({ name: "IMGLAB.AutoMask",

async beforeRegisterNodeDef(nodeType, nodeData, app) {
    if (nodeData.name !== "AutoMask") return;

    const origCreated = nodeType.prototype.onNodeCreated;
    nodeType.prototype.onNodeCreated = function () {
        origCreated?.apply(this, arguments);
        // Default node colour — purple theme
        if (!this.color)    this.color   = "#2d1a4a";  // header — dark purple
        if (!this.bgcolor)  this.bgcolor = "#4a2878";  // body — original purple
        this._am_matchResults = {};
        this._am_imageSize    = null;
        this._am_maskSizes    = {};
        this._am_elseMaskSize = null;
        for (let i = this.inputs.length - 1; i >= 0; i--) {
            const n = this.inputs[i].name;
            if (/^mask_[2-9]$/.test(n) || /^mask_\d{2,}$/.test(n) || n === "else")
                this.removeInput(i);
        }
        this._am_watchNoMatchWidget();
    };

    nodeType.prototype._am_watchNoMatchWidget = function () {
        const widget = this.widgets?.find(w => w.name === "no_match");
        if (!widget) return;
        const self = this;
        const orig = widget.callback;
        widget.callback = function (v) { orig?.call(this, v); self._am_syncElseSlot(v); };
        this._am_syncElseSlot(widget.value);
    };

    nodeType.prototype._am_syncElseSlot = function (value) {
        const hasElse = this.inputs.some(i => i.name === "else");
        if (value === "Use Else Mask" && !hasElse) {
            this.addInput("else", MASK_TYPE);
            this.setSize(this.computeSize());
        } else if (value !== "Use Else Mask" && hasElse) {
            const idx = this.inputs.findIndex(i => i.name === "else");
            if (idx >= 0) { this.removeInput(idx); this.setSize(this.computeSize()); }
        }
        this.setDirtyCanvas(true, true);
    };

    nodeType.prototype._am_getMaskInputs = function () {
        return this.inputs
            .map((inp, idx) => ({ ...inp, _idx: idx }))
            .filter(inp => /^mask_\d+$/.test(inp.name))
            .sort((a, b) => parseInt(a.name.split("_")[1]) - parseInt(b.name.split("_")[1]));
    };

    nodeType.prototype._am_updateDynamicInputs = function () {
        const masks = this._am_getMaskInputs();
        let lastConnected = -1;
        for (let i = 0; i < masks.length; i++) if (masks[i].link != null) lastConnected = i;
        const targetCount = Math.min(lastConnected + 2, MAX_MASKS);
        for (let n = masks.length + 1; n <= targetCount; n++) this.addInput(`mask_${n}`, MASK_TYPE);
        const updated = this._am_getMaskInputs();
        for (let i = updated.length - 1; i >= targetCount; i--) {
            const gIdx = this.inputs.findIndex(inp => inp.name === updated[i].name);
            if (gIdx >= 0 && updated[i].link == null) this.removeInput(gIdx);
        }
        // Keep else at the bottom
        const elseIdx = this.inputs.findIndex(inp => inp.name === "else");
        if (elseIdx >= 0 && elseIdx !== this.inputs.length - 1) {
            this.removeInput(elseIdx);
            this.addInput("else", MASK_TYPE);
        }
        this.setSize(this.computeSize());
    };

    nodeType.prototype.onConnectionsChange = function (type, slotIndex, connected, link_info, input) {
        if (type !== 1) return;
        this._am_updateDynamicInputs();
        this.setDirtyCanvas(true, true);
    };

    nodeType.prototype.onExecuted = function (message) {
        if (message?.match_results?.[0]) {
            const raw = message.match_results[0];
            this._am_matchResults = {};
            for (const k in raw) this._am_matchResults[parseInt(k)] = raw[k];
        }
        if (message?.image_size?.[0])   this._am_imageSize = message.image_size[0];
        if (message?.mask_sizes?.[0]) {
            const raw = message.mask_sizes[0];
            this._am_maskSizes = {};
            for (const k in raw) this._am_maskSizes[parseInt(k)] = raw[k];
        }
        this._am_elseMaskSize = message?.else_mask_size?.[0] ?? null;
        this.setDirtyCanvas(true, true);
    };

    nodeType.prototype._am_helpBtnRect = function () {
        const titleH = LiteGraph.NODE_TITLE_HEIGHT;
        return {
            x: this.size[0] - HELP_BTN_R - HELP_BTN_W,
            y: -titleH + (titleH - HELP_BTN_H) / 2,
            w: HELP_BTN_W, h: HELP_BTN_H,
        };
    };

    const origDrawFg = nodeType.prototype.onDrawForeground;
    nodeType.prototype.onDrawForeground = function (ctx) {
        origDrawFg?.apply(this, arguments);

        // ? button
        const btn = this._am_helpBtnRect();
        ctx.beginPath(); ctx.roundRect(btn.x, btn.y, btn.w, btn.h, 4);
        ctx.fillStyle = "rgba(255,255,255,0.12)"; ctx.fill();
        ctx.strokeStyle = "rgba(255,255,255,0.3)"; ctx.lineWidth = 1; ctx.stroke();
        ctx.font = "bold 11px sans-serif"; ctx.fillStyle = "rgba(255,255,255,0.75)";
        ctx.textAlign = "center";
        ctx.fillText("?", btn.x + btn.w / 2, btn.y + btn.h - 3);

        if (this.flags?.collapsed) return;

        const matchResults = this._am_matchResults || {};
        const maskSizes    = this._am_maskSizes    || {};
        const imageSize    = this._am_imageSize;
        const elseMaskSize = this._am_elseMaskSize;
        const nodeW        = this.size[0];

        // Compute shared dot column from widest mask_N name only.
        // image and else are shorter and would misalign the column.
        ctx.font = SLOT_FONT;
        let maxMaskNameW = ctx.measureText("mask_1").width; // minimum
        for (const inp of this.inputs) {
            if (/^mask_\d+$/.test(inp.name)) {
                const w = ctx.measureText(inp.name).width;
                if (w > maxMaskNameW) maxMaskNameW = w;
            }
        }
        const dotX = SLOT_NAME_X + maxMaskNameW + NAME_GAP + DOT_RADIUS;

        // Resolution text is right-aligned to just before the output label area.
        // Output labels ("image","mask","rgba") are at the right side of the node.
        // We reserve OUTPUT_LABEL_RESERVE px from the right for them.
        const resRight = nodeW - OUTPUT_LABEL_RESERVE;
		const dimX = dotX+DOT_GAP;
		
        const drawDot = (cx, cy, color, symbol) => {
            ctx.beginPath();
            ctx.arc(cx, cy, DOT_RADIUS, 0, Math.PI * 2);
            ctx.fillStyle = color; ctx.fill();
            if (symbol) {
                ctx.font = "bold 8px sans-serif"; ctx.fillStyle = "#fff";
                ctx.textAlign = "center";
                ctx.fillText(symbol, cx, cy + 3);
            }
        };

        const drawRes = (text, cy, color) => {
            ctx.font = "11px monospace"; ctx.fillStyle = color;
            ctx.textAlign = "left";
            ctx.fillText(text, dimX, cy + 4);
        };

        for (let i = 0; i < this.inputs.length; i++) {
            const inp       = this.inputs[i];
            const canvasPos = this.getConnectionPos(true, i);
            const cy        = canvasPos[1] - this.pos[1];

            // ── image ────────────────────────────────────────────────
            if (inp.name === "image") {
                if (imageSize) drawRes(`${imageSize[0]}×${imageSize[1]}`, cy, RES_COLOR);
                continue;
            }

            // ── else ─────────────────────────────────────────────────
            if (inp.name === "else") {
                ctx.font = SLOT_FONT;
                const elseNameW = ctx.measureText("else").width;
                const elseDotX  = SLOT_NAME_X + elseNameW + NAME_GAP + DOT_RADIUS;
                const connected = inp.link != null;
                drawDot(dotX, cy, connected ? ELSE_COLOR : NA_COLOR, connected && elseMaskSize ? "~" : null);
                if (connected && elseMaskSize) drawRes(`${elseMaskSize[0]}×${elseMaskSize[1]}`, cy, ELSE_COLOR);
                continue;
            }

            if (!/^mask_\d+$/.test(inp.name)) continue;

            const num       = parseInt(inp.name.split("_")[1]);
            const connected = inp.link != null;
            const hasResult = num in matchResults;
            const hasMaskSz = num in maskSizes;

            // dot at shared column
            if (connected && hasResult) {
                const isMatch = matchResults[num];
                drawDot(dotX, cy, isMatch ? OK_COLOR : ERR_COLOR, isMatch ? "✓" : "✗");
            } else if (!connected) {
                ctx.beginPath();
                ctx.arc(dotX, cy, DOT_RADIUS - 1, 0, Math.PI * 2);
                ctx.fillStyle = NA_COLOR; ctx.fill();
            }

            // resolution right-aligned before output labels
            if (connected && hasMaskSz) {
                const [mw, mh] = maskSizes[num];
                const isMatch  = hasResult ? matchResults[num] : null;
                drawRes(`${mw}×${mh}`, cy,
                    isMatch === true ? OK_COLOR : isMatch === false ? ERR_COLOR : RES_COLOR);
            }
        }
    };

    const origMouseDown = nodeType.prototype.onMouseDown;
    nodeType.prototype.onMouseDown = function (e, localPos, graphCanvas) {
        const btn = this._am_helpBtnRect();
        const [mx, my] = localPos;
        if (mx >= btn.x && mx <= btn.x + btn.w && my >= btn.y && my <= btn.y + btn.h) {
            const canvas = graphCanvas.canvas, rect = canvas.getBoundingClientRect();
            const scale = graphCanvas.ds?.scale ?? 1, offset = graphCanvas.ds?.offset ?? [0, 0];
            showHelpNear(
                rect.left + (this.pos[0] + offset[0]) * scale,
                rect.top  + (this.pos[1] + offset[1]) * scale,
                this.size[0] * scale
            );
            return true;
        }
        return origMouseDown?.call(this, e, localPos, graphCanvas);
    };
},

});

README.md Unescape Escape

README.md