import type { LGraphCanvas as TLGraphCanvas, LGraphNode, SerializedLGraphNode, serializedLGraph, ContextMenuItem, LGraph as TLGraph, AdjustedMouseEvent, IContextMenuOptions, } from "typings/litegraph.js"; import type {ComfyApiFormat, ComfyApiPrompt, ComfyApp} from "typings/comfy.js"; import {app} from "scripts/app.js"; import {api} from "scripts/api.js"; import {SERVICE as CONFIG_SERVICE} from "./services/config_service.js"; import {SERVICE as KEY_EVENT_SERVICE} from "./services/key_events_services.js"; import {fixBadLinks} from "rgthree/common/link_fixer.js"; import {injectCss, wait} from "rgthree/common/shared_utils.js"; import {replaceNode, waitForCanvas, waitForGraph} from "./utils.js"; import {NodeTypesString, addRgthree, getNodeTypeStrings, stripRgthree} from "./constants.js"; import {RgthreeProgressBar} from "rgthree/common/progress_bar.js"; import {RgthreeConfigDialog} from "./config.js"; import { iconGear, iconNode, iconReplace, iconStarFilled, logoRgthree, } from "rgthree/common/media/svgs.js"; import type {Bookmark} from "./bookmark.js"; import {createElement, query, queryOne} from "rgthree/common/utils_dom.js"; export enum LogLevel { IMPORTANT = 1, ERROR, WARN, INFO, DEBUG, DEV, } const LogLevelKeyToLogLevel: {[key: string]: LogLevel} = { IMPORTANT: LogLevel.IMPORTANT, ERROR: LogLevel.ERROR, WARN: LogLevel.WARN, INFO: LogLevel.INFO, DEBUG: LogLevel.DEBUG, DEV: LogLevel.DEV, }; 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]: "log", [LogLevel.DEV]: "log", }; const LogLevelToCSS: {[key in LogLevel]: string} = { [LogLevel.IMPORTANT]: "font-weight: bold; color: blue;", [LogLevel.ERROR]: "", [LogLevel.WARN]: "", [LogLevel.INFO]: "font-style: italic; color: blue;", [LogLevel.DEBUG]: "font-style: italic; color: #444;", [LogLevel.DEV]: "color: #004b68;", }; let GLOBAL_LOG_LEVEL = LogLevel.ERROR; /** * At some point in Summer of 2024 ComfyUI broke third-party api calls by assuming api paths follow * a certain structure. However, rgthree-comfy wants an `/rgthree/` prefix for that same reason, so * we overwrite the apiUrl method to fix. */ const apiURL = api.apiURL; api.apiURL = function (route: string): string { if (route.includes("rgthree/")) { return (this.api_base + "/" + route).replace(/\/\//g, "/"); } return apiURL.apply(this, arguments as any); }; /** * A blocklist of extensions to disallow hooking into rgthree's base classes when calling the * `rgthree.invokeExtensionsAsync` method (which runs outside of ComfyNode's * `app.invokeExtensionsAsync` which is private). * * In Apr 2024 the base rgthree node class added support for other extensions using `nodeCreated` * and `beforeRegisterNodeDef` which allows other extensions to modify the class. However, since it * had been months since divorcing the ComfyNode in rgthree-comfy due to instability and * inflexibility, this was a bit risky as other extensions hadn't ever run with this ability. This * list attempts to block extensions from being able to call into rgthree-comfy nodes via the * `nodeCreated` and `beforeRegisterNodeDef` callbacks now that rgthree-comfy is utilizing them * because they do not work. Oddly, it's ComfyUI's own extension that is broken. */ const INVOKE_EXTENSIONS_BLOCKLIST = [ { name: "Comfy.WidgetInputs", reason: "Major conflict with rgthree-comfy nodes' inputs causing instability and " + "repeated link disconnections.", }, { name: "efficiency.widgethider", reason: "Overrides value getter before widget getter is prepared. Can be lifted if/when " + "https://github.com/jags111/efficiency-nodes-comfyui/pull/203 is pulled.", }, ]; /** A basic wrapper around logger. */ class Logger { /** Logs a message to the console if it meets the current log level. */ log(level: LogLevel, message: string, ...args: any[]) { const [n, v] = this.logParts(level, message, ...args); console[n]?.(...v); } /** * Returns a tuple of the console function and its arguments. Useful for callers to make the * actual console. call to gain benefits of DevTools knowing the source line. * * If the input is invalid or the level doesn't meet the configuration level, then the return * value is an unknown function and empty set of values. Callers can use optionla chaining * successfully: * * const [fn, values] = logger.logPars(LogLevel.INFO, 'my message'); * console[fn]?.(...values); // Will work even if INFO won't be logged. * */ logParts(level: LogLevel, message: string, ...args: any[]): [ConsoleLogFns, any[]] { if (level <= GLOBAL_LOG_LEVEL) { const css = LogLevelToCSS[level] || ""; if (level === LogLevel.DEV) { message = `🔧 ${message}`; } return [LogLevelToMethod[level], [`%c${message}`, css, ...args]]; } return ["none" as "info", []]; } } /** * A log session, with the name as the prefix. A new session will stack prefixes. */ class LogSession { readonly logger = new Logger(); readonly logsCache: {[key: string]: {lastShownTime: number}} = {}; constructor(readonly name?: string) {} /** * Returns the console log method to use and the arguments to pass so the call site can log from * there. This extra work at the call site allows for easier debugging in the dev console. * * const [logMethod, logArgs] = logger.logParts(LogLevel.DEBUG, message, ...args); * console[logMethod]?.(...logArgs); */ logParts(level: LogLevel, message?: string, ...args: any[]): [ConsoleLogFns, any[]] { message = `${this.name || ""}${message ? " " + message : ""}`; return this.logger.logParts(level, message, ...args); } logPartsOnceForTime( level: LogLevel, time: number, message?: string, ...args: any[] ): [ConsoleLogFns, any[]] { message = `${this.name || ""}${message ? " " + message : ""}`; const cacheKey = `${level}:${message}`; const cacheEntry = this.logsCache[cacheKey]; const now = +new Date(); if (cacheEntry && cacheEntry.lastShownTime + time > now) { return ["none" as "info", []]; } const parts = this.logger.logParts(level, message, ...args); if (console[parts[0]]) { this.logsCache[cacheKey] = this.logsCache[cacheKey] || ({} as {lastShownTime: number}); this.logsCache[cacheKey]!.lastShownTime = now; } return parts; } debugParts(message?: string, ...args: any[]) { return this.logParts(LogLevel.DEBUG, message, ...args); } infoParts(message?: string, ...args: any[]) { return this.logParts(LogLevel.INFO, message, ...args); } warnParts(message?: string, ...args: any[]) { return this.logParts(LogLevel.WARN, message, ...args); } newSession(name?: string) { return new LogSession(`${this.name}${name}`); } } export type RgthreeUiMessage = { id: string; message: string; type?: "warn" | "info" | "success" | 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 { /** Exposes the ComfyUI api instance on rgthree. */ readonly api = api; private settingsDialog: RgthreeConfigDialog | null = null; private progressBarEl: RgthreeProgressBar | null = null; private rgthreeCssPromise: Promise; /** Stores a node id that we will use to queu only that output node (with `queueOutputNode`). */ private queueNodeIds: number[] | null = null; logger = new LogSession("[rgthree]"); monitorBadLinksAlerted = false; monitorLinkTimeout: number | null = null; processingQueue = false; loadingApiJson = false; replacingReroute: number | null = null; processingMouseDown = false; processingMouseUp = false; processingMouseMove = false; lastAdjustedMouseEvent: AdjustedMouseEvent | null = null; // Comfy/LiteGraph states so nodes and tell what the hell is going on. canvasCurrentlyCopyingToClipboard = false; canvasCurrentlyCopyingToClipboardWithMultipleNodes = false; initialGraphToPromptSerializedWorkflowBecauseComfyUIBrokeStuff: any = null; private readonly isMac: boolean = !!( navigator.platform?.toLocaleUpperCase().startsWith("MAC") || (navigator as any).userAgentData?.platform?.toLocaleUpperCase().startsWith("MAC") ); constructor() { super(); const logLevel = LogLevelKeyToLogLevel[CONFIG_SERVICE.getConfigValue("log_level")] ?? GLOBAL_LOG_LEVEL; this.setLogLevel(logLevel); this.initializeGraphAndCanvasHooks(); this.initializeComfyUIHooks(); this.initializeContextMenu(); this.rgthreeCssPromise = injectCss("extensions/rgthree-comfy/rgthree.css"); this.initializeProgressBar(); CONFIG_SERVICE.addEventListener("config-change", ((e: CustomEvent) => { if (e.detail?.key?.includes("features.progress_bar")) { this.initializeProgressBar(); } }) as EventListener); if (CONFIG_SERVICE.getConfigValue("debug.keys_down.enabled")) { const elDebugKeydowns = createElement("div.rgthree-debug-keydowns", { parent: document.body, }); const updateDebugKeyDown = () => { elDebugKeydowns.innerText = Object.keys(KEY_EVENT_SERVICE.downKeys).join(" "); } KEY_EVENT_SERVICE.addEventListener("keydown", updateDebugKeyDown); KEY_EVENT_SERVICE.addEventListener("keyup", updateDebugKeyDown); } } /** * Initializes the top progress bar, if it's configured. */ async initializeProgressBar() { if (CONFIG_SERVICE.getConfigValue("features.progress_bar.enabled")) { await this.rgthreeCssPromise; if (!this.progressBarEl) { this.progressBarEl = RgthreeProgressBar.create(); this.progressBarEl.setAttribute( "title", "Progress Bar by rgthree. right-click for rgthree menu.", ); this.progressBarEl.addEventListener("contextmenu", async (e) => { e.stopPropagation(); e.preventDefault(); }); this.progressBarEl.addEventListener("pointerdown", async (e) => { LiteGraph.closeAllContextMenus(); if (e.button == 2) { const canvas = await waitForCanvas(); new LiteGraph.ContextMenu( this.getRgthreeContextMenuItems(), { title: `
${logoRgthree} rgthree-comfy
`, left: e.clientX, top: 5, }, canvas.getCanvasWindow(), ); return; } if (e.button == 0) { const nodeId = this.progressBarEl?.currentNodeId; if (nodeId) { const [canvas, graph] = await Promise.all([waitForCanvas(), waitForGraph()]); const node = graph.getNodeById(Number(nodeId)); if (node) { canvas.centerOnNode(node); e.stopPropagation(); e.preventDefault(); } } return; } }); } // Handle both cases in case someone hasn't updated. Can probably just assume // `isUpdatedComfyBodyClasses` is true in the near future. const isUpdatedComfyBodyClasses = !!queryOne(".comfyui-body-top"); const position = CONFIG_SERVICE.getConfigValue("features.progress_bar.position"); this.progressBarEl.classList.toggle("rgthree-pos-bottom", position === "bottom"); // If ComfyUI is updated with the body segments, then use that. if (isUpdatedComfyBodyClasses) { if (position === "bottom") { queryOne(".comfyui-body-bottom")!.appendChild(this.progressBarEl); } else { queryOne(".comfyui-body-top")!.appendChild(this.progressBarEl); } } else { document.body.appendChild(this.progressBarEl); } const height = CONFIG_SERVICE.getConfigValue("features.progress_bar.height") || 14; this.progressBarEl.style.height = `${height}px`; const fontSize = Math.max(10, Number(height) - 10); this.progressBarEl.style.fontSize = `${fontSize}px`; this.progressBarEl.style.fontWeight = fontSize <= 12 ? "bold" : "normal"; } else { this.progressBarEl?.remove(); } } /** * Initialize a bunch of hooks into LiteGraph itself so we can either keep state or context on * what's happening so nodes can respond appropriately. This is usually to fix broken assumptions * in the unowned code [🤮], but sometimes to add features or enhancements too [⭐]. */ private async initializeGraphAndCanvasHooks() { const rgthree = this; // [🤮] 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 our nodes can find the seralized node. Works with method // `getNodeFromInitialGraphToPromptSerializedWorkflowBecauseComfyUIBrokeStuff` to find a node // while serializing. What a way to work around... const graphSerialize = LGraph.prototype.serialize; LGraph.prototype.serialize = function () { const response = graphSerialize.apply(this, [...arguments] as any) as any; rgthree.initialGraphToPromptSerializedWorkflowBecauseComfyUIBrokeStuff = response; return response; }; // Overrides LiteGraphs' processMouseDown to both keep state as well as dispatch a custom event. const processMouseDown = LGraphCanvas.prototype.processMouseDown; LGraphCanvas.prototype.processMouseDown = function (e: AdjustedMouseEvent) { rgthree.processingMouseDown = true; const returnVal = processMouseDown.apply(this, [...arguments] as any); rgthree.dispatchCustomEvent("on-process-mouse-down", {originalEvent: e}); rgthree.processingMouseDown = false; return returnVal; }; // Overrides LiteGraph's `adjustMouseEvent` to capture the last even coming in and out. Useful // to capture the last `canvasX` and `canvasY` properties, which are not the same as LiteGraph's // `canvas.last_mouse_position`, unfortunately. const adjustMouseEvent = LGraphCanvas.prototype.adjustMouseEvent; LGraphCanvas.prototype.adjustMouseEvent = function (e: PointerEvent) { adjustMouseEvent.apply(this, [...arguments] as any); rgthree.lastAdjustedMouseEvent = e as AdjustedMouseEvent; }; // [🤮] Copying to clipboard clones nodes and then manipulats the linking data manually which // does not allow a node to handle connections. This harms nodes that manually handle inputs, // like our any-input nodes that may start with one input, and manually add new ones when one is // attached. const copyToClipboard = LGraphCanvas.prototype.copyToClipboard; LGraphCanvas.prototype.copyToClipboard = function (nodes: LGraphNode[]) { rgthree.canvasCurrentlyCopyingToClipboard = true; rgthree.canvasCurrentlyCopyingToClipboardWithMultipleNodes = Object.values(nodes || this.selected_nodes || []).length > 1; copyToClipboard.apply(this, [...arguments] as any); rgthree.canvasCurrentlyCopyingToClipboard = false; rgthree.canvasCurrentlyCopyingToClipboardWithMultipleNodes = false; }; // [⭐] Make it so when we add a group, we get to name it immediately. const onGroupAdd = LGraphCanvas.onGroupAdd; LGraphCanvas.onGroupAdd = function (...args: any[]) { const graph = app.graph as TLGraph; onGroupAdd.apply(this, [...args] as any); LGraphCanvas.onShowPropertyEditor( {}, null, null, null, graph._groups[graph._groups.length - 1], ); }; } /** * [🤮] Handles the same exact thing as ComfyApp's `invokeExtensionsAsync`, but done here since * it is #private in ComfyApp because... of course it us. This is necessary since we purposefully * avoid using the ComfyNode due to historical instability and inflexibility for all the advanced * ui stuff rgthree-comfy nodes do, but we can still have other custom nodes know what's happening * with rgthree-comfy; specifically, for `nodeCreated` as of now. */ async invokeExtensionsAsync(method: "nodeCreated", ...args: any[]) { const comfyapp = app as ComfyApp; if (CONFIG_SERVICE.getConfigValue("features.invoke_extensions_async.node_created") === false) { const [m, a] = this.logParts( LogLevel.INFO, `Skipping invokeExtensionsAsync for applicable rgthree-comfy nodes`, ); console[m]?.(...a); return Promise.resolve(); } return await Promise.all( comfyapp.extensions.map(async (ext) => { if (ext?.[method]) { try { const blocked = INVOKE_EXTENSIONS_BLOCKLIST.find((block) => ext.name.toLowerCase().startsWith(block.name.toLowerCase()), ); if (blocked) { const [n, v] = this.logger.logPartsOnceForTime( LogLevel.WARN, 5000, `Blocked extension '${ext.name}' method '${method}' for rgthree-nodes because: ${blocked.reason}`, ); console[n]?.(...v); return Promise.resolve(); } return await (ext[method] as Function)(...args, comfyapp); } catch (error) { const [n, v] = this.logParts( LogLevel.ERROR, `Error calling extension '${ext.name}' method '${method}' for rgthree-node.`, {error}, {extension: ext}, {args}, ); console[n]?.(...v); } } }), ); } /** * Wraps `dispatchEvent` for easier CustomEvent dispatching. */ private dispatchCustomEvent(event: string, detail?: any) { if (detail != null) { return this.dispatchEvent(new CustomEvent(event, {detail})); } return this.dispatchEvent(new CustomEvent(event)); } /** * Initializes hooks specific to an rgthree-comfy context menu on the root menu. */ private async initializeContextMenu() { const that = this; setTimeout(async () => { const getCanvasMenuOptions = LGraphCanvas.prototype.getCanvasMenuOptions; LGraphCanvas.prototype.getCanvasMenuOptions = function (...args: any[]) { let existingOptions = getCanvasMenuOptions.apply(this, [...args] as any); const options = []; options.push(null); // Divider options.push(null); // Divider options.push(null); // Divider options.push({ content: logoRgthree + `rgthree-comfy`, className: "rgthree-contextmenu-item rgthree-contextmenu-main-item-rgthree-comfy", submenu: { options: that.getRgthreeContextMenuItems(), }, }); options.push(null); // Divider options.push(null); // Divider let idx = null; idx = idx || existingOptions.findIndex((o) => o?.content?.startsWith?.("Queue Group")) + 1; idx = idx || existingOptions.findIndex((o) => o?.content?.startsWith?.("Queue Selected")) + 1; idx = idx || existingOptions.findIndex((o) => o?.content?.startsWith?.("Convert to Group")); idx = idx || existingOptions.findIndex((o) => o?.content?.startsWith?.("Arrange (")); idx = idx || existingOptions.findIndex((o) => !o) + 1; idx = idx || 3; existingOptions.splice(idx, 0, ...options); for (let i = existingOptions.length; i > 0; i--) { if (existingOptions[i] === null && existingOptions[i + 1] === null) { existingOptions.splice(i, 1); } } return existingOptions; }; }, 1016); } /** * Returns the standard menu items for an rgthree-comfy context menu. */ private getRgthreeContextMenuItems(): ContextMenuItem[] { const [canvas, graph] = [app.canvas as TLGraphCanvas, app.graph as TLGraph]; const selectedNodes = Object.values(canvas.selected_nodes || {}); let rerouteNodes: LGraphNode[] = []; if (selectedNodes.length) { rerouteNodes = selectedNodes.filter((n) => n.type === "Reroute"); } else { rerouteNodes = graph._nodes.filter((n) => n.type == "Reroute"); } const rerouteLabel = selectedNodes.length ? "selected" : "all"; const showBookmarks = CONFIG_SERVICE.getFeatureValue("menu_bookmarks.enabled"); const bookmarkMenuItems = showBookmarks ? getBookmarks() : []; return [ { content: "Nodes", disabled: true, className: "rgthree-contextmenu-item rgthree-contextmenu-label", }, { content: iconNode + "All", className: "rgthree-contextmenu-item", has_submenu: true, submenu: { options: getNodeTypeStrings() as unknown as ContextMenuItem[], callback: ( value: string | ContextMenuItem, options: IContextMenuOptions, event: MouseEvent, ) => { const node = LiteGraph.createNode(addRgthree(value as string)); node.pos = [ rgthree.lastAdjustedMouseEvent!.canvasX, rgthree.lastAdjustedMouseEvent!.canvasY, ]; canvas.graph.add(node); canvas.selectNode(node); app.graph.setDirtyCanvas(true, true); }, extra: {rgthree_doNotNest: true}, }, }, { content: "Actions", disabled: true, className: "rgthree-contextmenu-item rgthree-contextmenu-label", }, { content: iconGear + "Settings (rgthree-comfy)", disabled: !!this.settingsDialog, className: "rgthree-contextmenu-item", callback: (...args: any[]) => { this.settingsDialog = new RgthreeConfigDialog().show(); this.settingsDialog.addEventListener("close", (e) => { this.settingsDialog = null; }); }, }, { content: iconReplace + ` Convert ${rerouteLabel} Reroutes`, disabled: !rerouteNodes.length, className: "rgthree-contextmenu-item", callback: (...args: any[]) => { const msg = `Convert ${rerouteLabel} ComfyUI Reroutes to Reroute (rgthree) nodes? \n` + `(First save a copy of your workflow & check reroute connections afterwards)`; if (!window.confirm(msg)) { return; } (async () => { for (const node of [...rerouteNodes]) { if (node.type == "Reroute") { this.replacingReroute = node.id; await replaceNode(node, NodeTypesString.REROUTE); this.replacingReroute = null; } } })(); }, }, ...bookmarkMenuItems, { content: "More...", disabled: true, className: "rgthree-contextmenu-item rgthree-contextmenu-label", }, { content: iconStarFilled + "Star on Github", className: "rgthree-contextmenu-item rgthree-contextmenu-github", callback: (...args: any[]) => { window.open("https://github.com/rgthree/rgthree-comfy", "_blank"); }, }, ]; } /** * Wraps an `app.queuePrompt` call setting a specific node id that we will inspect and change the * serialized graph right before being sent (below, in our `api.queuePrompt` override). */ async queueOutputNodes(nodeIds: number[]) { try { this.queueNodeIds = nodeIds; await app.queuePrompt(); } catch (e) { const [n, v] = this.logParts( LogLevel.ERROR, `There was an error queuing nodes ${nodeIds}`, e, ); console[n]?.(...v); } finally { this.queueNodeIds = null; } } /** * Recusively walks backwards from a node adding its inputs to the `newOutput` from `oldOutput`. */ private recursiveAddNodes(nodeId: string, oldOutput: ComfyApiFormat, newOutput: ComfyApiFormat) { let currentId = nodeId; let currentNode = oldOutput[currentId]!; if (newOutput[currentId] == null) { newOutput[currentId] = currentNode; for (const inputValue of Object.values(currentNode.inputs || [])) { if (Array.isArray(inputValue)) { this.recursiveAddNodes(inputValue[0], oldOutput, newOutput); } } } return newOutput; } /** * Initialize a bunch of hooks into ComfyUI and/or LiteGraph itself so we can either keep state or * context on what's happening so nodes can respond appropriately. This is usually to fix broken * assumptions in the unowned code [🤮], but sometimes to add features or enhancements too [⭐]. */ private initializeComfyUIHooks() { const rgthree = this; // Keep state for when the app is queuing the prompt. For instance, this is used for seed to // understand if we're serializing because we're queueing (and return the random seed to use) or // for saving the workflow (and keep -1, etc.). const queuePrompt = app.queuePrompt as Function; app.queuePrompt = async function () { rgthree.processingQueue = true; rgthree.dispatchCustomEvent("queue"); try { await queuePrompt.apply(app, [...arguments]); } finally { rgthree.processingQueue = false; rgthree.dispatchCustomEvent("queue-end"); } }; // Keep state for when the app is in the middle of loading from an api JSON file. const loadApiJson = app.loadApiJson; app.loadApiJson = async function () { rgthree.loadingApiJson = true; try { loadApiJson.apply(app, [...arguments] as any); } finally { rgthree.loadingApiJson = false; } }; // Keep state for when the app is serizalizing the graph to prompt. const graphToPrompt = app.graphToPrompt; app.graphToPrompt = async function () { rgthree.dispatchCustomEvent("graph-to-prompt"); let promise = graphToPrompt.apply(app, [...arguments] as any); await promise; rgthree.dispatchCustomEvent("graph-to-prompt-end"); return promise; }; // Override the queuePrompt for api to intercept the prompt output and, if queueNodeIds is set, // then we only want to queue those nodes, by rewriting the api format (prompt 'output' field) // so only those are evaluated. const apiQueuePrompt = api.queuePrompt as Function; api.queuePrompt = async function (index: number, prompt: ComfyApiPrompt) { if (rgthree.queueNodeIds?.length && prompt.output) { const oldOutput = prompt.output; let newOutput = {}; for (const queueNodeId of rgthree.queueNodeIds) { rgthree.recursiveAddNodes(String(queueNodeId), oldOutput, newOutput); } prompt.output = newOutput; } rgthree.dispatchCustomEvent("comfy-api-queue-prompt-before", { workflow: prompt.workflow, output: prompt.output, }); const response = apiQueuePrompt.apply(app, [index, prompt]); rgthree.dispatchCustomEvent("comfy-api-queue-prompt-end"); return response; }; // Hook into a clean call; allow us to clear and rgthree messages. const clean = app.clean; app.clean = function () { rgthree.clearAllMessages(); clean && clean.apply(app, [...arguments] as any); }; // Hook into a data load, like from an image or JSON drop-in. This is (currently) used to // monitor for bad linking data. const loadGraphData = app.loadGraphData; app.loadGraphData = function (graph: serializedLGraph) { if (rgthree.monitorLinkTimeout) { clearTimeout(rgthree.monitorLinkTimeout); rgthree.monitorLinkTimeout = null; } rgthree.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 as unknown as TLGraph); if (fixBadLinksResult.hasBadLinks) { const [n, v] = rgthree.logParts( LogLevel.WARN, `The workflow you've loaded has corrupt linking data. Open ${ new URL(location.href).origin }/rgthree/link_fixer to try to fix.`, ); console[n]?.(...v); if (CONFIG_SERVICE.getConfigValue("features.show_alerts_for_corrupt_workflows")) { rgthree.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: "/rgthree/link_fixer", }, { label: "Fix in place", href: "/rgthree/link_fixer", 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 as unknown as TLGraph, true, ); if (!fixBadLinksResult.hasBadLinks) { rgthree.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 ( CONFIG_SERVICE.getConfigValue("features.monitor_for_corrupt_links") || CONFIG_SERVICE.getConfigValue("features.monitor_bad_links") ) { rgthree.monitorLinkTimeout = setTimeout(() => { rgthree.monitorBadLinks(); }, 5000); } } } catch (e) { console.error(e); alert("Unsuccessful at fixing corrupt data. :("); rgthree.hideMessage("bad-links"); } } }, }, ], }); } } else if ( CONFIG_SERVICE.getConfigValue("features.monitor_for_corrupt_links") || CONFIG_SERVICE.getConfigValue("features.monitor_bad_links") ) { rgthree.monitorLinkTimeout = setTimeout(() => { rgthree.monitorBadLinks(); }, 5000); } }, 100); return loadGraphData && loadGraphData.apply(app, [...arguments] as any); }; } /** * [🤮] Finds a node in the currently serializing workflow from the hook setup above. This is to * mitigate breakages from https://github.com/comfyanonymous/ComfyUI/issues/2193 we can try to * store the workflow node so our nodes can find the seralized node. */ getNodeFromInitialGraphToPromptSerializedWorkflowBecauseComfyUIBrokeStuff( node: LGraphNode, ): SerializedLGraphNode | null { return ( this.initialGraphToPromptSerializedWorkflowBecauseComfyUIBrokeStuff?.nodes?.find( (n: SerializedLGraphNode) => n.id === node.id, ) ?? null ); } /** * Shows a message in the UI. */ 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); } // If we have a dialog open then we want to append the message to the dialog so they show over // the modal. const dialogs = query("dialog[open]"); if (dialogs.length) { let dialog = dialogs[dialogs.length - 1]!; dialog.appendChild(container); dialog.addEventListener("close", (e) => { 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.innerHTML = 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 = " | "; 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); } } /** * Hides a message in the UI. */ 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(); } /** * Clears all messages in the UI. */ async clearAllMessages() { let container = document.querySelector(".rgthree-top-messages-container"); container && (container.innerHTML = ""); } setLogLevel(level?: LogLevel | string) { if (typeof level === "string") { level = LogLevelKeyToLogLevel[CONFIG_SERVICE.getConfigValue("log_level")]; } if (level != null) { GLOBAL_LOG_LEVEL = level; } } logParts(level: LogLevel, message?: string, ...args: any[]) { return this.logger.logParts(level, message, ...args); } newLogSession(name?: string) { return this.logger.newSession(name); } isDebugMode() { if (window.location.href.includes("rgthree-debug=false")) { return false; } return GLOBAL_LOG_LEVEL >= LogLevel.DEBUG || window.location.href.includes("rgthree-debug"); } isDevMode() { if (window.location.href.includes("rgthree-dev=false")) { return false; } return GLOBAL_LOG_LEVEL >= LogLevel.DEV || window.location.href.includes("rgthree-dev"); } 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); } } function getBookmarks(): ContextMenuItem[] { const graph: TLGraph = app.graph; // Sorts by Title. // I could see an option to sort by either Shortcut, Title, or Position. const bookmarks = graph._nodes .filter((n): n is Bookmark => n.type === NodeTypesString.BOOKMARK) .sort((a, b) => a.title.localeCompare(b.title)) .map((n) => ({ content: `[${n.shortcutKey}] ${n.title}`, className: "rgthree-contextmenu-item", callback: () => { n.canvasToBookmark(); }, })); return !bookmarks.length ? [] : [ { content: "🔖 Bookmarks", disabled: true, className: "rgthree-contextmenu-item rgthree-contextmenu-label", }, ...bookmarks, ]; } export const rgthree = new Rgthree(); // Expose it on window because, why not. (window as any).rgthree = rgthree;