cg-use-everywhere/js/use_everywhere.js

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

import { is_UEnode, is_helper, inject, Logger, get_real_node, defineProperty } from "./use_everywhere_utilities.js";
import { displayMessage, update_input_label, indicate_restriction, UpdateBlocker } from "./use_everywhere_ui.js";
import { LinkRenderController } from "./use_everywhere_ui.js";
import { autoCreateMenu } from "./use_everywhere_autocreate.js";
import { add_autoprompts } from "./use_everywhere_autoprompt.js";
import { GraphAnalyser } from "./use_everywhere_graph_analysis.js";
import { main_menu_settings, node_menu_settings, canvas_menu_settings, non_ue_menu_settings } from "./use_everywhere_settings.js";
import { add_debug } from "./ue_debug.js";

/*
The ui component that looks after the link rendering
*/
var linkRenderController;
var graphAnalyser;

/*
Inject a call to linkRenderController.mark_list_link_outdated into a method with name methodname on all objects in the array
If object is undefined, do nothing.
The injection is added at the end of the existing method (if the method didn't exist, it is created).
A Logger.trace call is added at the start with 'tracetext'
*/
function inject_outdating_into_objects(array, methodname, tracetext) {
    if (array) {
        array.forEach((object) => { inject_outdating_into_object_method(object, methodname, tracetext); })
    }
}
function inject_outdating_into_object_method(object, methodname, tracetext) {
    if (object) inject(object, methodname, tracetext, linkRenderController.mark_link_list_outdated, linkRenderController);
}

app.registerExtension({
	name: "cg.customnodes.use_everywhere",

    async beforeRegisterNodeDef(nodeType, nodeData, app) {
        /*
        When a node is connected or unconnected, the link list is dirty.
        If it is a UE node, we need to update it as well
        */
        const onConnectionsChange = nodeType.prototype.onConnectionsChange;
        nodeType.prototype.onConnectionsChange = function (side,slot,connect,link_info,output) {        
            Logger.trace("onConnectionsChange", arguments, this);
            if (this.IS_UE && side==1) { // side 1 is input
                if (this.type=="Anything Everywhere?" && slot!=0) {
                    // don't do anything for the regexs
                } else {
                    const type = (connect && link_info) ? get_real_node(link_info?.origin_id)?.outputs[link_info?.origin_slot]?.type : undefined;
                    this.input_type[slot] = type;
                    if (link_info) link_info.type = type ? type : "*";
                    update_input_label(this, slot, app);
                }
            }
            linkRenderController.mark_link_list_outdated();
            onConnectionsChange?.apply(this, arguments);
        };

        /*
        Extra menu options are the node right click menu.
        We add to this list, and also insert a link list outdate to everything.
        */
        const getExtraMenuOptions = nodeType.prototype.getExtraMenuOptions;
        nodeType.prototype.getExtraMenuOptions = function(_, options) {
            Logger.trace("getExtraMenuOptions", arguments, this);
            getExtraMenuOptions?.apply(this, arguments);
            if (is_UEnode(this)) {
                node_menu_settings(options, this);
            } else {
                non_ue_menu_settings(options, this);
            }
            inject_outdating_into_objects(options,'callback',`menu option on ${this.id}`);
        }

        /*
        When a UE node is created, we set the group and color restriction properties.
        We also create pseudo-widgets for all the inputs so that they can be searched
        and to avoid other code throwing errors.
        */
        if (is_UEnode(nodeType)) {
            const onNodeCreated = nodeType.prototype.onNodeCreated;
            nodeType.prototype.onNodeCreated = function () {
                const r = onNodeCreated ? onNodeCreated.apply(this, arguments) : undefined;
                if (!this.properties) this.properties = {}
                this.properties.group_restricted = 0;
                this.properties.color_restricted = 0;
                if (this.inputs) {
                    if (!this.widgets) this.widgets = [];
                    for (const input of this.inputs) {
                        if (input.widget && !this.widgets.find((w) => w.name === input.widget.name)) this.widgets.push(input.widget)
                    }
                }
                return r;
            }
        }
    },

    async nodeCreated(node) {
        if (!node.__mode) {
            node.__mode = node.mode
            defineProperty(node, "mode", {
                get: ( )=>{return node.__mode},
                set: (v)=>{node.__mode = v; node.afterChangeMade?.('mode', v);}            
            })
        }
        if (!node.__bgcolor) {
            node.__bgcolor = node.bgcolor
            defineProperty(node,"bgcolor", {
                get: ( )=>{return node.__bgcolor},
                set: (v)=>{node.__bgcolor = v; node.afterChangeMade?.('bgcolor', v);}                       
            })
        }
        const acm = node.afterChangeMade
        node.afterChangeMade = (p, v) => {
            acm?.(p,v)
            if (p==='bgcolor') {
                if (node.mode!=4) linkRenderController.mark_link_list_outdated();
            }
            if (p==='mode') {
                linkRenderController.mark_link_list_outdated();
                node.widgets?.forEach((widget) => {widget.onModeChange?.(v)});
            }
        }

        node.IS_UE = is_UEnode(node);
        if (node.IS_UE) {
            node.input_type = [undefined, undefined, undefined]; // for dynamic input types
            node.displayMessage = displayMessage;                // receive messages from the python code           

            // If a widget on a UE node is edited, link list is dirty
            inject_outdating_into_objects(node.widgets,'callback',`widget callback on ${node.id}`);

            // draw the indication of group restrictions
            const original_onDrawTitleBar = node.onDrawTitleBar;
            node.onDrawTitleBar = function(ctx, title_height) {
                original_onDrawTitleBar?.apply(this, arguments);
                if (node.properties.group_restricted || node.properties.color_restricted) indicate_restriction(ctx, title_height);
            }
        }

        if (is_helper(node)) { // editing a helper node makes the list dirty
            inject_outdating_into_objects(node.widgets,'callback',`widget callback on ${this.id}`);
        }

        // removing a node makes the list dirty
        inject_outdating_into_object_method(node, 'onRemoved', `node ${node.id} removed`)

        // creating a node makes the link list dirty - but give the system a moment to finish
        setTimeout( ()=>{linkRenderController.mark_link_list_outdated()}, 100 );
    }, 

    // When a graph node is loaded collapsed the UI need to know
    // probably not needed now autocomplete is gone?
    loadedGraphNode(node) { if (node.flags.collapsed && node.loaded_when_collapsed) node.loaded_when_collapsed(); },

	async setup() {
        /*
        Add css for the autocomplete. Probably not needed now
        */
        const head = document.getElementsByTagName('HEAD')[0];
        const link = document.createElement('link');
        link.rel = 'stylesheet';
        link.type = 'text/css';
        link.href = 'extensions/cg-use-everywhere/ue.css';
        head.appendChild(link);

        /*
        Listen for message-handler event from python code
        */
        function messageHandler(event) {
            const id = event.detail.id;
            const message = event.detail.message;
            const node = get_real_node(id);
            if (node && node.displayMessage) node.displayMessage(id, message);
            else (console.log(`node ${id} couldn't handle a message`));
        }
        api.addEventListener("ue-message-handler", messageHandler);

        api.addEventListener("status", ({detail}) => {
            if (linkRenderController) linkRenderController.note_queue_size(detail ? detail.exec_info.queue_remaining : 0)
        });

        /*
        Don't modify the graph when saving the workflow or api
        */
        const _original_save_onclick = document.getElementById('comfy-save-button').onclick;
        document.getElementById('comfy-save-button').onclick = function() {
            graphAnalyser.pause();
            _original_save_onclick();
            graphAnalyser.unpause()
        }
        const _original_save_api_onclick = document.getElementById('comfy-dev-save-api-button').onclick;
        document.getElementById('comfy-dev-save-api-button').onclick = function() {
            graphAnalyser.pause();
            // should check for UE links here and give a warning: #217
            _original_save_api_onclick();
            graphAnalyser.unpause();
        }
        
        /* 
        When we draw a node, render the virtual connection points
        */
        const original_drawNode = LGraphCanvas.prototype.drawNode;
        LGraphCanvas.prototype.drawNode = function(node, ctx) {
            UpdateBlocker.push()
            try {
                const v = original_drawNode.apply(this, arguments);
                linkRenderController.highlight_ue_connections(node, ctx);
                return v
            } finally { UpdateBlocker.pop() }
        }

        /*
        When we draw connections, do the ue ones as well (logic for on/off is in lrc)
        */
        const drawConnections = LGraphCanvas.prototype.drawConnections;
        LGraphCanvas.prototype.drawConnections = function(ctx) {
            drawConnections?.apply(this, arguments);
            linkRenderController.render_all_ue_links(ctx);
        }

        /*
        Add to the main settings
        */
        main_menu_settings();
        
        /* 
        Canvas menu is the right click on backdrop.
        We need to add our options, and hijack the others to mark link list dirty
        */
        const original_getCanvasMenuOptions = LGraphCanvas.prototype.getCanvasMenuOptions;
        LGraphCanvas.prototype.getCanvasMenuOptions = function () {
            // Add our items to the canvas menu 
            const options = original_getCanvasMenuOptions.apply(this, arguments);
            canvas_menu_settings(options);
            
            //  every menu item makes our list dirty
            inject_outdating_into_objects(options,'callback',`menu option on canvas`);

            return options;
        }

        /*
        When you drag from a node, showConnectionMenu is called. If shift key is pressed call ours
        Broken #219
        */
        const showSearchBox = LGraphCanvas.prototype.showSearchBox;
        LGraphCanvas.prototype.showSearchBox = function (optPass) {
            if (optPass.shiftKey) {
                autoCreateMenu.apply(this, arguments);
            } else {
                this.use_original_menu = true;
                showSearchBox.apply(this, arguments);
                this.use_original_menu = false;
            }
        }

        /*
        To allow us to use the shift drag above, we need to intercept 'allow_searchbox' sometimes
        (because searchbox is the default behaviour when shift dragging)
        Broken #219
        */
        var original_allow_searchbox = app.canvas.allow_searchbox;
        defineProperty(app.canvas, 'allow_searchbox', {
            get : function() { 
                if (this.use_original_menu) { return original_allow_searchbox; }
                if(app.ui.settings.getSettingValue('AE.replacesearch', true) && this.connecting_output) {
                    return false;
                } else { return original_allow_searchbox; }
            },
            set : function(v) { original_allow_searchbox = v; }
        });
        

	},

    init() {
        graphAnalyser = GraphAnalyser.instance();
        app.graphToPrompt = async function () {
            return graphAnalyser.analyse_graph(true, true, false);
        }
        
        linkRenderController = LinkRenderController.instance(graphAnalyser);

        add_autoprompts();

        if (false) add_debug();

    }

});