rgthree-rgthree-comfy/ts/rgthree.ts

import type { LGraphNode, SerializedLGraphNode, serializedLGraph } from "litegraph.js";
// @ts-ignore
import { app } from "../../scripts/app.js";
// @ts-ignore
import { rgthreeConfig } from "./rgthree_config.js";
import { fixBadLinks } from "./link_fixer.js";
import { wait } from "./shared_utils.js";
import { waitForGraph } from "./utils.js";

export enum LogLevel {
  IMPORTANT = 1,
  ERROR,
  WARN,
  INFO,
  DEBUG,
}

type ConsoleLogFns = "log" | "error" | "warn" | "debug" | "info";
const LogLevelToMethod: { [key in LogLevel]: ConsoleLogFns } = {
  [LogLevel.IMPORTANT]: "log",
  [LogLevel.ERROR]: "error",
  [LogLevel.WARN]: "warn",
  [LogLevel.INFO]: "info",
  [LogLevel.DEBUG]: "debug",
};
const LogLevelToCSS: { [key in LogLevel]: string } = {
  [LogLevel.IMPORTANT]: "font-weight:bold; color:blue;",
  [LogLevel.ERROR]: "",
  [LogLevel.WARN]: "",
  [LogLevel.INFO]: "",
  [LogLevel.DEBUG]: "font-style: italic;",
};

let GLOBAL_LOG_LEVEL = LogLevel.DEBUG;

/** A basic wrapper around logger. */
class Logger {
  log(level: LogLevel, message: string, ...args: any[]) {
    if (level <= GLOBAL_LOG_LEVEL) {
      const css = LogLevelToCSS[level] || "";
      console[LogLevelToMethod[level]](`%c${message}`, css, ...args);
    }
  }
}

/**
 * A log session, with the name as the prefix. A new session will stack prefixes.
 */
class LogSession {
  logger = new Logger();
  constructor(readonly name?: string) {}

  log(levelOrMessage: LogLevel | string, message?: string, ...args: any[]) {
    let level = typeof levelOrMessage === "string" ? LogLevel.INFO : levelOrMessage;
    if (typeof levelOrMessage === "string") {
      message = levelOrMessage;
    }
    this.logger.log(level, `${this.name || ""}${message ? " " + message : ""}`, ...args);
  }

  debug(message?: string, ...args: any[]) {
    this.log(LogLevel.DEBUG, message, ...args);
  }

  info(message?: string, ...args: any[]) {
    this.log(LogLevel.INFO, message, ...args);
  }

  error(message?: string, ...args: any[]) {
    this.log(LogLevel.ERROR, message, ...args);
  }

  newSession(name?: string) {
    return new LogSession(`${this.name}${name}`);
  }
}

export type RgthreeUiMessage = {
  id: string;
  message: string;
  type?: "warn" | null;
  timeout?: number;
  // closeable?: boolean; // TODO
  actions?: Array<{
    label: string;
    href?: string;
    callback?: (event: MouseEvent) => void;
  }>;
};

/**
 * A global class as 'rgthree'; exposed on wiindow. Lots can go in here.
 */
class Rgthree extends EventTarget {
  /** Are any functional keys pressed in this given moment? */
  ctrlKey = false;
  altKey = false;
  metaKey = false;
  shiftKey = false;
  readonly downKeys: { [key: string]: boolean } = {};

  logger = new LogSession("[rgthree]");

  monitorBadLinksAlerted = false;
  monitorLinkTimeout: number | null = null;

  processingQueue = false;

  initialGraphToPromptSerializedWorkflowBecauseComfyUIBrokeStuff: any = null;

  async clearAllMessages() {
    let container = document.querySelector(".rgthree-top-messages-container");
    container && (container.innerHTML = "");
  }

  async showMessage(data: RgthreeUiMessage) {
    let container = document.querySelector(".rgthree-top-messages-container");
    if (!container) {
      container = document.createElement("div");
      container.classList.add("rgthree-top-messages-container");
      document.body.appendChild(container);
    }
    // Hide if we exist.
    await this.hideMessage(data.id);

    const messageContainer = document.createElement("div");
    messageContainer.setAttribute("type", data.type || "info");

    const message = document.createElement("span");
    message.innerText = data.message;
    messageContainer.appendChild(message);

    for (let a = 0; a < (data.actions || []).length; a++) {
      const action = data.actions![a]!;
      if (a > 0) {
        const sep = document.createElement("span");
        sep.innerHTML = "&nbsp;|&nbsp;";
        messageContainer.appendChild(sep);
      }

      const actionEl = document.createElement("a");
      actionEl.innerText = action.label;
      if (action.href) {
        actionEl.target = "_blank";
        actionEl.href = action.href;
      }
      if (action.callback) {
        actionEl.onclick = (e) => {
          return action.callback!(e);
        };
      }
      messageContainer.appendChild(actionEl);
    }

    const messageAnimContainer = document.createElement("div");
    messageAnimContainer.setAttribute("msg-id", data.id);
    messageAnimContainer.appendChild(messageContainer);
    container.appendChild(messageAnimContainer);

    // Add. Wait. Measure. Wait. Anim.
    await wait(64);
    messageAnimContainer.style.marginTop = `-${messageAnimContainer.offsetHeight}px`;
    await wait(64);
    messageAnimContainer.classList.add("-show");

    if (data.timeout) {
      await wait(data.timeout);
      this.hideMessage(data.id);
    }
  }

  async hideMessage(id: string) {
    const msg = document.querySelector(`.rgthree-top-messages-container > [msg-id="${id}"]`);
    if (msg?.classList.contains("-show")) {
      msg.classList.remove("-show");
      await wait(750);
    }
    msg && msg.remove();
  }

  /** Handle keydown. Pulled out because sometimes a node will get a keydown before rgthree. */
  handleKeydown(e: KeyboardEvent) {
    this.ctrlKey = !!e.ctrlKey;
    this.altKey = !!e.altKey;
    this.metaKey = !!e.metaKey;
    this.shiftKey = !!e.shiftKey;
    this.downKeys[e.key.toLocaleUpperCase()] = true;
    this.downKeys["^" + e.key.toLocaleUpperCase()] = true;
  }

  /** Handle keyup. Pulled out because sometimes a node will get a keyup before rgthree. */
  handleKeyup(e: KeyboardEvent) {
    this.ctrlKey = !!e.ctrlKey;
    this.altKey = !!e.altKey;
    this.metaKey = !!e.metaKey;
    this.shiftKey = !!e.shiftKey;
    this.downKeys[e.key.toLocaleUpperCase()] = false;
    this.downKeys["^" + e.key.toLocaleUpperCase()] = false;
  }

  areAllKeysDown(keys: string[], caseSensitive = false) {
    return keys.every((k) => {
      if (caseSensitive) {
        return rgthree.downKeys["^" + k.trim()];
      }
      return rgthree.downKeys[k.trim().toUpperCase()];
    });
  }

  // This sucks.. but to mitigate changes from https://github.com/rgthree/rgthree-comfy/issues/69
  // and https://github.com/comfyanonymous/ComfyUI/issues/2193 we can try to store the workflow
  // node so out nodes can find the seralized node. What a way to work around...
  getNodeFromInitialGraphToPromptSerializedWorkflowBecauseComfyUIBrokeStuff(node: LGraphNode) : SerializedLGraphNode|null {
    return this.initialGraphToPromptSerializedWorkflowBecauseComfyUIBrokeStuff?.nodes?.find(
      (n: SerializedLGraphNode) => n.id === node.id,
    ) ?? null;
  }


  private async initializeGraphHooks() {
    const graph = await waitForGraph();
    const onSerialize = (graph as any).onSerialize;
    (graph as any).onSerialize = (data: any) => {
      // This sucks.. but to mitigate changes from https://github.com/rgthree/rgthree-comfy/issues/69
      // and https://github.com/comfyanonymous/ComfyUI/issues/2193 we can try to store the workflow
      // node so out nodes can find the seralized node. What a way to work around...
      this.initialGraphToPromptSerializedWorkflowBecauseComfyUIBrokeStuff = data;
      onSerialize?.apply(graph, data);
    }
  }

  constructor() {
    super();
    window.addEventListener("keydown", (e) => {
      this.handleKeydown(e);
    });

    window.addEventListener("keyup", (e) => {
      this.handleKeyup(e);
    });

    this.initializeGraphHooks();

    // Override the loadGraphData so we can check for bad links and ask the user to fix them.
    const that = this;

    const queuePrompt = app.queuePrompt as Function;
    app.queuePrompt = async function () {
      that.dispatchEvent(new CustomEvent("queue"));
      that.processingQueue = true;
      try {
        await queuePrompt.apply(app, [...arguments]);
      } finally {
        that.processingQueue = false;
        that.dispatchEvent(new CustomEvent("queue-end"));
      }
    };

    const graphToPrompt = app.graphToPrompt as Function;
    app.graphToPrompt = async function () {
      that.dispatchEvent(new CustomEvent("graph-to-prompt"));
      let promise = graphToPrompt.apply(app, [...arguments]);
      await promise;
      that.dispatchEvent(new CustomEvent("graph-to-prompt-end"));
      return promise;
    };

    const clean = app.clean;
    app.clean = function () {
      that.clearAllMessages();
      clean && clean.call(app, ...arguments);
    };

    const loadGraphData = app.loadGraphData;
    app.loadGraphData = function (graph: serializedLGraph) {
      if (that.monitorLinkTimeout) {
        clearTimeout(that.monitorLinkTimeout);
        that.monitorLinkTimeout = null;
      }
      that.clearAllMessages();
      // Try to make a copy to use, because ComfyUI's loadGraphData will modify it.
      let graphCopy: serializedLGraph | null;
      try {
        graphCopy = JSON.parse(JSON.stringify(graph));
      } catch (e) {
        graphCopy = null;
      }
      setTimeout(() => {
        const wasLoadingAborted = document
          .querySelector(".comfy-modal-content")
          ?.textContent?.includes("Loading aborted due");
        const graphToUse = wasLoadingAborted ? graphCopy || graph : app.graph;
        const fixBadLinksResult = fixBadLinks(graphToUse);
        if (fixBadLinksResult.hasBadLinks) {
          that.log(
            LogLevel.WARN,
            `The workflow you've loaded has corrupt linking data. Open ${
              new URL(location.href).origin
            }/extensions/rgthree-comfy/html/links.html to try to fix.`,
          );
          if (rgthreeConfig["show_alerts_for_corrupt_workflows"]) {
            that.showMessage({
              id: "bad-links",
              type: "warn",
              message:
                "The workflow you've loaded has corrupt linking data that may be able to be fixed.",
              actions: [
                {
                  label: "Open fixer",
                  href: "/extensions/rgthree-comfy/html/links.html",
                },
                {
                  label: "Fix in place",
                  href: "/extensions/rgthree-comfy/html/links.html",
                  callback: (event) => {
                    event.stopPropagation();
                    event.preventDefault();
                    if (
                      confirm(
                        "This will attempt to fix in place. Please make sure to have a saved copy of your workflow.",
                      )
                    ) {
                      try {
                        const fixBadLinksResult = fixBadLinks(graphToUse, true);
                        if (!fixBadLinksResult.hasBadLinks) {
                          that.hideMessage("bad-links");
                          alert(
                            "Success! It's possible some valid links may have been affected. Please check and verify your workflow.",
                          );
                          wasLoadingAborted && app.loadGraphData(fixBadLinksResult.graph);
                          if (
                            rgthreeConfig["monitor_for_corrupt_links"] ||
                            rgthreeConfig["monitor_bad_links"]
                          ) {
                            that.monitorLinkTimeout = setTimeout(() => {
                              that.monitorBadLinks();
                            }, 5000);
                          }
                        }
                      } catch (e) {
                        console.error(e);
                        alert("Unsuccessful at fixing corrupt data. :(");
                        that.hideMessage("bad-links");
                      }
                    }
                  },
                },
              ],
            });
          }
        } else if (
          rgthreeConfig["monitor_for_corrupt_links"] ||
          rgthreeConfig["monitor_bad_links"]
        ) {
          that.monitorLinkTimeout = setTimeout(() => {
            that.monitorBadLinks();
          }, 5000);
        }
      }, 100);
      loadGraphData && loadGraphData.call(app, ...arguments);
    };

    wait(100).then(() => {
      this.injectRgthreeCss();
    });
  }

  private injectRgthreeCss() {
    let link = document.createElement("link");
    link.rel = "stylesheet";
    link.type = "text/css";
    link.href = "extensions/rgthree-comfy/rgthree.css";
    document.head.appendChild(link);
  }

  setLogLevel(level: LogLevel) {
    GLOBAL_LOG_LEVEL = level;
  }

  log(levelOrMessage: LogLevel | string, message?: string, ...args: any[]) {
    this.logger.log(levelOrMessage, message, ...args);
  }

  newLogSession(name?: string) {
    return this.logger.newSession(name);
  }

  monitorBadLinks() {
    const badLinksFound = fixBadLinks(app.graph);
    if (badLinksFound.hasBadLinks && !this.monitorBadLinksAlerted) {
      this.monitorBadLinksAlerted = true;
      alert(
        `Problematic links just found in live data. Can you save your workflow and file a bug with `
        + `the last few steps you took to trigger this at `
        + `https://github.com/rgthree/rgthree-comfy/issues. Thank you!`,
      );
    } else if (!badLinksFound.hasBadLinks) {
      // Clear the alert once fixed so we can alert again.
      this.monitorBadLinksAlerted = false;
    }
    this.monitorLinkTimeout = setTimeout(() => {
      this.monitorBadLinks();
    }, 5000);
  }
}

export const rgthree = new Rgthree();
// @ts-ignore. Expose it on window because, why not.
window.rgthree = rgthree;