chou18194766xx-comfyui_EncryptPreview/web/encrypt_preview.js

import { app } from "../../scripts/app.js";
import { api } from "../../scripts/api.js";
import { PreviewLoader } from "./loader.mjs";
import { NativePreview } from "./native_preview.mjs";

const loader = new PreviewLoader((route, options) => api.fetchApi(route, options));
const views = new Map();
let activePromptId = null;
let visibilityTimer;

const nativeCanvas = () => app.extensionManager.setting.get("Comfy.VueNodes.Enabled") !== true;
const cachedResults = node => node.graph === app.rootGraph ? app.nodeOutputs[String(node.id)]?.encrypt_preview : undefined;

function checkVisibility() {
  const enabled = nativeCanvas();
  const visibleNodes = new Set(app.canvas?.visible_nodes || []);
  for (const [node, view] of views) {
    view.setVisible(enabled && !node.flags.collapsed && visibleNodes.has(node));
    view.status();
  }
}

function removeView(node) {
  views.get(node)?.dispose();
  views.delete(node);
  if (!views.size) { clearInterval(visibilityTimer); visibilityTimer = null; }
}

api.addEventListener("execution_start", ({ detail }) => { activePromptId = detail.prompt_id; });
api.addEventListener("encrypt-preview", ({ detail }) => {
  if (!activePromptId || detail.prompt_id !== activePromptId) return;
  for (const [node, view] of views) {
    if (String(node.id) === detail.node_id && node.graph === app.rootGraph) {
      if (detail.event === "start") view.lastResults = cachedResults(node);
      view.receive(detail);
    }
  }
});

app.registerExtension({
  name: "Comfy.EncryptPreview",
  beforeRegisterNodeDef(nodeType, nodeData) {
    if (nodeData.name !== "EncryptPreviewImage") return;
    const created = nodeType.prototype.onNodeCreated;
    nodeType.prototype.onNodeCreated = function () {
      const result = created?.apply(this, arguments);
      const status = document.createElement("div");
      status.className = "encrypt-preview-status";
      Object.assign(status.style, { fontSize: "12px", lineHeight: "24px", padding: "0 4px", whiteSpace: "nowrap", overflow: "hidden", textOverflow: "ellipsis" });
      const widget = this.addDOMWidget("encrypt_preview_status", "div", status, {
        serialize: false, hideOnZoom: false, margin: 0, getMinHeight: () => 24, getMaxHeight: () => 24,
        getValue: () => "", setValue: () => {},
      });
      widget.serialize = false;
      widget.computeLayoutSize = () => ({ minHeight: 24, maxHeight: 24, minWidth: 100 });
      const view = new NativePreview(this, loader, (text, error) => {
        const supported = nativeCanvas();
        const message = supported ? text : "请关闭 Nodes 2.0 以显示预览";
        const state = supported && error ? "error" : "info";
        if (status.textContent === message && status.dataset.state === state) return;
        status.textContent = message;
        status.dataset.state = state;
        status.title = status.textContent;
        status.style.color = state === "error" ? "#ef7070" : "var(--input-text, #ddd)";
        this.graph?.setDirtyCanvas(true);
      });
      view.status();
      views.set(this, view);
      this.setSize([Math.max(this.size[0], 320), Math.max(this.size[1], 320)]);
      visibilityTimer ??= setInterval(checkVisibility, 500);
      return result;
    };

    const draw = nodeType.prototype.onDrawBackground;
    nodeType.prototype.onDrawBackground = function () {
      const view = views.get(this);
      const results = cachedResults(this);
      // ComfyUI can restore history/workflow outputs without an executed event.
      if (view && results && results !== view.lastResults) {
        view.lastResults = results;
        view.showResults(results);
      }
      if (!this.flags.collapsed && nativeCanvas()) view?.setVisible(true);
      return draw?.apply(this, arguments);
    };

    const executed = nodeType.prototype.onExecuted;
    nodeType.prototype.onExecuted = function (message) {
      const result = executed?.apply(this, arguments);
      const view = views.get(this);
      if (view && message?.encrypt_preview) {
        view.lastResults = message.encrypt_preview;
        view.showResults(message.encrypt_preview);
      }
      return result;
    };

    const menu = nodeType.prototype.getExtraMenuOptions;
    nodeType.prototype.getExtraMenuOptions = function (canvas, options) {
      const result = menu?.apply(this, arguments);
      const view = views.get(this);
      if (!view?.total) return result;
      const index = this.imageIndex ?? this.overIndex ?? 0;
      // The native actions would use the thumbnail while the original is loading.
      for (let i = options.length - 1; i >= 0; i--) {
        if (["Open Image", "Save Image", "Copy Image", "Copy (Clipspace)", "Paste (Clipspace)", "Open in MaskEditor | Image Canvas"].includes(options[i]?.content)) options.splice(i, 1);
      }
      options.unshift(
        { content: "查看原图", callback: () => { view.select(index); view.reconcile(); } },
        { content: "保存原图", callback: () => view.save(index) },
        { content: "重试预览", callback: () => view.retry() },
      );
      return result;
    };

    const removed = nodeType.prototype.onRemoved;
    nodeType.prototype.onRemoved = function () { removeView(this); return removed?.apply(this, arguments); };
  },
  beforeConfigureGraph() {
    activePromptId = null;
    for (const node of [...views.keys()]) removeView(node);
  },
});