Moser9815-ComfyUI-RMAutomation/js/rm_image_comparer.js

import { app } from "../../scripts/app.js";
import { api } from "../../scripts/api.js";

/**
 * RM Image Comparer.
 *
 * A/B wipe between two images. Written from scratch rather than lifted from rgthree, whose
 * version depends on its own base-node and widget framework — importing that would mean
 * installing most of the pack to get one node.
 *
 * Drawing happens in onDrawBackground against the node's own canvas rather than through DOM
 * overlays, so the comparison pans and zooms with the graph instead of floating above it.
 *
 * Interaction: hovering wipes, clicking locks the slider so you can look away without losing
 * the position, and clicking again releases it. The locked state is serialised with the
 * workflow because a comparison you set up deliberately should survive a reload.
 */

const NODE_NAME = "RMImageComparer";

/** Gap under the widgets before the image area starts. */
const TOP_PAD = 4;
const HANDLE_W = 2;
const HANDLE_GRAB = 12;

function imageURL(entry) {
  if (!entry) return null;
  const params = new URLSearchParams({
    filename: entry.filename ?? "",
    subfolder: entry.subfolder ?? "",
    type: entry.type ?? "temp",
  });
  // Previews reuse filenames across runs; without a cache-buster the browser happily shows
  // the previous execution's image.
  params.set("rand", String(Math.random()));
  return api.apiURL(`/view?${params.toString()}`);
}

/** Fit `natural` inside `box` preserving aspect ratio, centred. */
function fitRect(naturalW, naturalH, boxX, boxY, boxW, boxH) {
  if (!naturalW || !naturalH || boxW <= 0 || boxH <= 0) return null;
  const scale = Math.min(boxW / naturalW, boxH / naturalH);
  const w = naturalW * scale;
  const h = naturalH * scale;
  return { x: boxX + (boxW - w) / 2, y: boxY + (boxH - h) / 2, w, h };
}

app.registerExtension({
  name: "RMAutomation.ImageComparer",

  async beforeRegisterNodeDef(nodeType, nodeData) {
    if (nodeData.name !== NODE_NAME) return;

    const onNodeCreated = nodeType.prototype.onNodeCreated;
    nodeType.prototype.onNodeCreated = function () {
      const r = onNodeCreated?.apply(this, arguments);

      this.rmCompare = {
        a: [],           // loaded HTMLImageElements for side A
        b: [],
        index: 0,        // which frame of the batch to show
        slider: 0.5,     // 0..1 across the drawn image rect
        locked: false,
        rect: null,      // last drawn image rect, for hit-testing
      };
      this.rmIndexWidget = null;

      if (this.size[0] < 240) this.size[0] = 240;
      if (this.size[1] < 260) this.size[1] = 260;

      return r;
    };

    /** Load previews returned by the Python half. */
    nodeType.prototype.rmSetImages = function (aEntries, bEntries) {
      const state = this.rmCompare;
      const load = (entries) =>
        (entries ?? []).map((entry) => {
          const img = new Image();
          // Redraw as each image arrives; otherwise the node stays blank until the whole
          // batch has decoded.
          img.onload = () => this.setDirtyCanvas(true, true);
          img.src = imageURL(entry);
          return img;
        });

      state.a = load(aEntries);
      state.b = load(bEntries);
      state.index = 0;

      this.rmSyncIndexWidget();
      this.setDirtyCanvas(true, true);
    };

    /**
     * Show an index widget only while there is a batch to step through. Adding it in JS
     * rather than declaring it in Python matters: a Python widget is an input, so changing
     * it would invalidate the cache and re-run the graph just to look at another frame.
     */
    nodeType.prototype.rmSyncIndexWidget = function () {
      const state = this.rmCompare;
      const count = Math.max(state.a.length, state.b.length);

      if (count <= 1) {
        if (this.rmIndexWidget) {
          const i = this.widgets.indexOf(this.rmIndexWidget);
          if (i > -1) this.widgets.splice(i, 1);
          this.rmIndexWidget = null;
        }
        return;
      }

      if (!this.rmIndexWidget) {
        this.rmIndexWidget = this.addWidget(
          "number",
          "index",
          0,
          (v) => {
            const max = Math.max(this.rmCompare.a.length, this.rmCompare.b.length) - 1;
            this.rmCompare.index = Math.max(0, Math.min(max, Math.round(v)));
            this.rmIndexWidget.value = this.rmCompare.index;
            this.setDirtyCanvas(true, true);
          },
          { min: 0, max: count - 1, step: 10, precision: 0 },
        );
        // Widgets are not inputs; keep them out of the serialised prompt.
        this.rmIndexWidget.serialize = false;
      }
      this.rmIndexWidget.options.max = count - 1;
      this.rmCompare.index = Math.min(this.rmCompare.index, count - 1);
      this.rmIndexWidget.value = this.rmCompare.index;
    };

    const onExecuted = nodeType.prototype.onExecuted;
    nodeType.prototype.onExecuted = function (message) {
      onExecuted?.apply(this, arguments);
      this.rmSetImages(message?.a_images, message?.b_images);
    };

    const onDrawBackground = nodeType.prototype.onDrawBackground;
    nodeType.prototype.onDrawBackground = function (ctx) {
      onDrawBackground?.apply(this, arguments);
      if (this.flags?.collapsed) return;

      const state = this.rmCompare;
      if (!state) return;

      const top = this.widgets_start_y != null
        ? this.widgets_start_y + (this.widgets?.length ?? 0) * 24 + TOP_PAD
        : TOP_PAD;
      const boxX = 8;
      const boxY = top;
      const boxW = this.size[0] - 16;
      const boxH = this.size[1] - top - 8;
      if (boxW <= 0 || boxH <= 0) return;

      const imgA = state.a[Math.min(state.index, state.a.length - 1)];
      const imgB = state.b[Math.min(state.index, state.b.length - 1)];
      const ready = (img) => img && img.naturalWidth > 0;

      if (!ready(imgA) && !ready(imgB)) {
        state.rect = null;
        ctx.save();
        ctx.fillStyle = "#666";
        ctx.font = "12px sans-serif";
        ctx.textAlign = "center";
        ctx.fillText("Run the graph to compare", boxX + boxW / 2, boxY + boxH / 2);
        ctx.restore();
        return;
      }

      // Size the frame from whichever side we have, so a single wired input still previews.
      const ref = ready(imgA) ? imgA : imgB;
      const rect = fitRect(ref.naturalWidth, ref.naturalHeight, boxX, boxY, boxW, boxH);
      if (!rect) return;
      state.rect = rect;

      ctx.save();

      if (ready(imgA)) ctx.drawImage(imgA, rect.x, rect.y, rect.w, rect.h);

      // Only wipe when there are genuinely two sides; with one image the slider would just
      // erase half the preview for no reason.
      if (ready(imgB) && ready(imgA)) {
        const splitX = rect.x + rect.w * state.slider;
        ctx.save();
        ctx.beginPath();
        ctx.rect(splitX, rect.y, rect.x + rect.w - splitX, rect.h);
        ctx.clip();
        ctx.drawImage(imgB, rect.x, rect.y, rect.w, rect.h);
        ctx.restore();

        ctx.fillStyle = state.locked ? "#ffb300" : "#ffffff";
        ctx.fillRect(splitX - HANDLE_W / 2, rect.y, HANDLE_W, rect.h);
      } else if (ready(imgB)) {
        ctx.drawImage(imgB, rect.x, rect.y, rect.w, rect.h);
      }

      ctx.restore();
    };

    /** Map a node-local point onto the slider track. Returns null when outside the image. */
    nodeType.prototype.rmSliderFromPos = function (pos) {
      const rect = this.rmCompare?.rect;
      if (!rect) return null;
      const [x, y] = pos;
      if (y < rect.y || y > rect.y + rect.h) return null;
      return Math.max(0, Math.min(1, (x - rect.x) / rect.w));
    };

    const onMouseMove = nodeType.prototype.onMouseMove;
    nodeType.prototype.onMouseMove = function (event, pos, canvas) {
      onMouseMove?.apply(this, arguments);
      const state = this.rmCompare;
      if (!state || state.locked) return;
      const t = this.rmSliderFromPos(pos);
      if (t == null) return;
      state.slider = t;
      this.setDirtyCanvas(true, true);
    };

    const onMouseDown = nodeType.prototype.onMouseDown;
    nodeType.prototype.onMouseDown = function (event, pos, canvas) {
      const handled = onMouseDown?.apply(this, arguments);
      if (handled) return handled;

      const state = this.rmCompare;
      if (!state?.rect) return false;
      const t = this.rmSliderFromPos(pos);
      if (t == null) return false;

      // Clicking near a locked handle grabs it; clicking anywhere else toggles the lock.
      const splitX = state.rect.x + state.rect.w * state.slider;
      const nearHandle = Math.abs(pos[0] - splitX) <= HANDLE_GRAB;

      if (state.locked && nearHandle) {
        state.slider = t;
      } else {
        state.locked = !state.locked;
        if (state.locked) state.slider = t;
      }
      this.setDirtyCanvas(true, true);
      return true;
    };

    const onMouseLeave = nodeType.prototype.onMouseLeave;
    nodeType.prototype.onMouseLeave = function () {
      onMouseLeave?.apply(this, arguments);
      const state = this.rmCompare;
      if (!state || state.locked) return;
      // Park it centred so the node reads as a comparison at a glance when not hovered.
      state.slider = 0.5;
      this.setDirtyCanvas(true, true);
    };

    const onSerialize = nodeType.prototype.onSerialize;
    nodeType.prototype.onSerialize = function (info) {
      onSerialize?.apply(this, arguments);
      const state = this.rmCompare;
      if (!state) return;
      info.rm_compare = { slider: state.slider, locked: state.locked };
    };

    const onConfigure = nodeType.prototype.onConfigure;
    nodeType.prototype.onConfigure = function (info) {
      onConfigure?.apply(this, arguments);
      const saved = info?.rm_compare;
      if (!saved || !this.rmCompare) return;
      if (typeof saved.slider === "number") this.rmCompare.slider = saved.slider;
      this.rmCompare.locked = !!saved.locked;
    };
  },
});