1st implementation

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 bmad4ever
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,2 +1,19 @@
-# ComfyUI-Bmad-DirtyUndoRedo
+# Bmad's Dirty Undo-Redo for [ComfyUI](https://github.com/comfyanonymous/ComfyUI)  
+
 ComfyUI extension that adds undo (and redo) functionality.
+
+## Instalation:
+
+- Navigate to `/ComfyUI/custom_nodes/` folder;
+- `git clone git clone https://github.com/bmad4ever/ComfyUI-Bmad-DirtyUndoRedo`.
+
+
+### Troubleshooting
+
+This extension overrides some prototypes methods, but should keep the original behavior intact. Make sure this extension is the last to be loaded so that it affects all node types (web extensions appear to be loaded in alphabetic order with respect to their path). 
+
+
+### Why "dirty"?
+
+I was mainly concerned with getting a somewhat working solution quickly and not so much with implementing it "right". For example, the undo/redo data contains snapshots of the entire serialized workflow (this is overkill), and I'm not considering changing it.
+

__init__.py

@@ -0,0 +1,34 @@
+import os
+#import subprocess
+#import importlib.util
+import sys
+import filecmp
+import shutil
+
+import __main__
+
+python = sys.executable
+
+
+extentions_folder = os.path.join(os.path.dirname(os.path.realpath(__main__.__file__)),
+                                 "web" + os.sep + "extensions" + os.sep + "ZZZ-Bmad-DirtyUndoRedo")
+javascript_folder = os.path.join(os.path.dirname(os.path.realpath(__file__)), "js")
+
+
+if not os.path.exists(extentions_folder):
+    os.mkdir(extentions_folder)
+
+result = filecmp.dircmp(javascript_folder, extentions_folder)
+
+if result.left_only or result.diff_files:
+    file_list = list(result.left_only)
+    file_list.extend(x for x in result.diff_files if x not in file_list)
+
+    for file in file_list:
+        src_file = os.path.join(javascript_folder, file)
+        dst_file = os.path.join(extentions_folder, file)
+        if os.path.exists(dst_file):
+            os.remove(dst_file)
+        shutil.copy(src_file, dst_file)
+
+print('\033[92mBmad-DirtyUndoRedo Loaded.\033[0m')

js/DirtyUndoRedo.js

@@ -0,0 +1,210 @@
+import { app } from "/scripts/app.js";
+import { workflowHistory } from "/extensions/ZZZ-Bmad-DirtyUndoRedo/WorkflowHistory.js";
+
+//IMPORTANT, this extension should be the last to be loaded, so make sure the folder is the last alphabetically
+
+app.registerExtension({
+    name: "Comfy.Bmad.DirtyUndoRedo",
+    async beforeRegisterNodeDef(nodeType, nodeData, app) {
+
+        const origGetExtraMenuOptions = nodeType.prototype.getExtraMenuOptions;
+        nodeType.prototype.getExtraMenuOptions = function (_, options) {
+            const r = origGetExtraMenuOptions ? origGetExtraMenuOptions.apply(this, arguments) : undefined;
+            //console.log(options);
+            if (options !== undefined)
+                for (let i = 0; i < options.length; i++) if (options[i]) {
+                    const opt = options[i];
+                    if (opt.has_submenu) continue;
+                    const o_callback = opt.callback;
+                    opt.callback = function () {
+                        workflowHistory.prev_undo_timestamp -= workflowHistory.state_merge_threshold;
+                        workflowHistory.before();
+                        workflowHistory.lock();
+                        let ocr = undefined;
+                        try {
+                            ocr = o_callback ? o_callback.apply(this, arguments) : undefined;
+                        } catch (e) { }
+                        workflowHistory.release();
+                        workflowHistory.after();
+                        return ocr;
+                    }
+                }
+            return r;
+        }
+
+        const onNodeCreated = nodeType.prototype.onNodeCreated;
+        nodeType.prototype.onNodeCreated = function () {
+            const r = onNodeCreated ? onNodeCreated.apply(this, arguments) : undefined;
+
+            this.onMouseDown = function (event) {
+                //console.log("onMouseDown");
+
+                // set fixed node order so that state comparison works properly and avoids repeating the same state
+                // (is this really required tho? why not use _nodes_in_order when serializing? )
+                for (let i = 0; i < app.graph._nodes.length; i++) app.graph._nodes[i] = app.graph._nodes_in_order[i];
+                return false;
+            };
+
+            return r;
+        }
+
+    },
+    loadedGraphNode(node, _) {
+        if (app.graph._nodes_in_order[0] !== node) return;// execute only once
+
+        workflowHistory.setup(app);
+
+        // do not clean undo history if loading was triggered by an undo/redo command
+        if (app.wh.disable_load_reset) 
+        {
+            app.wh.disable_load_reset = false;
+            return;
+        }
+
+        app.wh.clean_history(app);
+        app.wh.get_new_candidate_state(app);
+    }
+});
+
+
+//=======================================================================================
+//          Override LiteGraph default behaviors
+
+LGraphCanvas.onMenuNodeColors = function (value, options, e, menu, node) {
+    if (!node) {
+        throw "no node for color";
+    }
+
+    var values = [];
+    values.push({
+        value: null,
+        content:
+            "<span style='display: block; padding-left: 4px;'>No color</span>"
+    });
+
+    for (var i in LGraphCanvas.node_colors) {
+        var color = LGraphCanvas.node_colors[i];
+        var value = {
+            value: i,
+            content:
+                "<span style='display: block; color: #999; padding-left: 4px; border-left: 8px solid " +
+                color.color +
+                "; background-color:" +
+                color.bgcolor +
+                "'>" +
+                i +
+                "</span>"
+        };
+        values.push(value);
+    }
+    new LiteGraph.ContextMenu(values, {
+        event: e,
+        callback: inner_clicked,
+        parentMenu: menu,
+        node: node
+    });
+
+    function inner_clicked(v) {
+        if (!node) {
+            return;
+        }
+
+        var color = v.value ? LGraphCanvas.node_colors[v.value] : null;
+
+        var fApplyColor = function (node) {
+            if (color) {
+                if (node.constructor === LiteGraph.LGraphGroup) {
+                    node.color = color.groupcolor;
+                } else {
+                    node.color = color.color;
+                    node.bgcolor = color.bgcolor;
+                }
+            } else {
+                delete node.color;
+                delete node.bgcolor;
+            }
+        }
+
+        workflowHistory.prev_undo_timestamp -= workflowHistory.state_merge_threshold; // force potential push state
+        workflowHistory.before();
+        workflowHistory.lock();
+        var graphcanvas = LGraphCanvas.active_canvas;
+        if (!graphcanvas.selected_nodes || Object.keys(graphcanvas.selected_nodes).length <= 1) {
+            fApplyColor(node);
+        } else {
+            for (var i in graphcanvas.selected_nodes) {
+                fApplyColor(graphcanvas.selected_nodes[i]);
+            }
+        }
+        node.setDirtyCanvas(true, true);
+        workflowHistory.release();
+        workflowHistory.after();
+    }
+
+    return false;
+};
+
+
+LGraphCanvas.onMenuNodeShapes = function (value, options, e, menu, node) {
+    if (!node) {
+        throw "no node passed";
+    }
+
+    new LiteGraph.ContextMenu(LiteGraph.VALID_SHAPES, {
+        event: e,
+        callback: inner_clicked,
+        parentMenu: menu,
+        node: node
+    });
+
+    function inner_clicked(v) {
+        if (!node) {
+            return;
+        }
+        node.graph.beforeChange(/*?*/); //node
+
+        var fApplyMultiNode = function (node) {
+            node.shape = v;
+        }
+
+        workflowHistory.prev_undo_timestamp -= workflowHistory.state_merge_threshold; // force potential push state
+        workflowHistory.before();
+        workflowHistory.lock();
+        var graphcanvas = LGraphCanvas.active_canvas;
+        if (!graphcanvas.selected_nodes || Object.keys(graphcanvas.selected_nodes).length <= 1) {
+            fApplyMultiNode(node);
+        } else {
+            for (var i in graphcanvas.selected_nodes) {
+                fApplyMultiNode(graphcanvas.selected_nodes[i]);
+            }
+        }
+
+        node.graph.afterChange(/*?*/); //node
+        node.setDirtyCanvas(true);
+        workflowHistory.release();
+        workflowHistory.after();
+    }
+
+    return false;
+}
+
+
+LGraphCanvas.prototype.original_prompt = LGraphCanvas.prototype.prompt;
+LGraphCanvas.prototype.prompt = function(title, value, callback, event, multiline) {
+    //console.log("LGraphCanvas.prototype.prompt");
+    const new_callback = function(node){
+        //workflowHistory.before() should have been triggered on mouse click
+        callback(node);
+        workflowHistory.after();
+    }
+    return this.original_prompt(title, value, new_callback, event, multiline);
+};
+
+
+LiteGraph.ContextMenu.prototype.original_close = LiteGraph.ContextMenu.prototype.close;
+LiteGraph.ContextMenu.prototype.close = function(e, ignore_parent_menu) {
+    //console.log("LiteGraph.ContextMenu.prototype.close");
+    const r = this.original_close(e, ignore_parent_menu);
+    workflowHistory.after();
+    return r;
+}

js/WorkflowHistory.js

@@ -0,0 +1,251 @@
+class WorkflowHistory {
+
+    // STORES ENTIRE GRAPH STATE *NOT THE CHANGES !!!
+
+
+    lock() { this.enabled = false; }
+    release() { this.enabled = true; }
+    isBetweenOnBeforeAndOnAfter() { return this.b_count > 0; }
+
+    constructor() {
+        // the following paramaters CAN BE TWEAKED
+        this.max_undo_steps = 100; // the maximum number of undos that can be performed
+        this.max_redo_steps = 50; // the maximum number of redos that can be performed
+        this.state_merge_threshold = 100; // state merge threshold in milliseconds
+
+        // DO NOT CHANGE THE VARIABLES BELOW
+        //WorkflowHistory.WH = this;
+        this.undo_history = [];
+        this.redo_history = [];
+        this.prev_undo_timestamp = 0; // timestamp of the last state pushed to the undo history
+        this.setup_done = false;
+
+        // auxiliary variables used by the nodes (to avoid data spam)
+        this.temp_timestamp = 0;   //potential state to push to undo history
+        this.temp_state = 0;
+
+        this.disable_load_reset = false; //distinguish when a load is triggered by an undo/redo command, or by loading a new workflow
+        this.between_onbefore_onafter = false; //ignores conection changes when true
+        this.b_count = 0;
+
+        this.enabled = true;
+
+
+        this.setup = function (app) {
+            if (this.setup_done) return;
+            this.setup_done = true;
+            app.wh = this;
+
+            const o_onMouse = app.canvas.onMouse;
+            app.canvas.onMouse = function () {
+                const r = o_onMouse ? o_onMouse.apply(this, arguments) : undefined;
+
+                if (!workflowHistory.enabled) return r;
+                //console.log("onMouse");
+                workflowHistory.before();
+
+                return r;
+            }
+
+            const o__mouseup_callback = app.canvas._mouseup_callback;
+            app.canvas._mouseup_callback = function (event) {
+                const r = o__mouseup_callback ? o__mouseup_callback.apply(this, arguments) : undefined;
+
+                if (!workflowHistory.enabled) return r;
+                //console.log("_mouseup_callback");
+
+                const leftClick = 1;
+                const rightClick = 3;
+
+                if (event.which === leftClick) {
+                    workflowHistory.after();
+                }
+
+                if (event.which === rightClick) {
+                    workflowHistory.before(); // ??
+                }
+
+                return r;
+            };
+
+            //app.graph.onConnectionChange // not needed apparently
+
+            this.repKeyCount = 0;
+            window.addEventListener("keydown", function (event) {
+                if(workflowHistory.prevKey === event.key) return; //avoid quick auto spamm
+                else if(workflowHistory.prevKey !== null) {
+                    clearTimeout(workflowHistory.keyTimeout);
+                    workflowHistory.repKeyCount=0;
+                }
+                
+                if (event.ctrlKey) {
+                    if (event.key === "z") {
+                        workflowHistory.undo(app);
+                        workflowHistory.prevKey = "z";
+                        workflowHistory.keyTimeout = setTimeout(
+                            ()=>{workflowHistory.prevKey=null}, 
+                            workflowHistory.time_to_next_operation_repeat() 
+                            );
+                        workflowHistory.repKeyCount += 1;
+                    }
+                    if (event.key === "y") {
+                        workflowHistory.redo(app);
+                        workflowHistory.prevKey = "y";
+                        workflowHistory.keyTimeout = setTimeout(
+                            ()=>{workflowHistory.prevKey=null}, 
+                            workflowHistory.time_to_next_operation_repeat() 
+                            );
+                        workflowHistory.repKeyCount += 1;
+                    }
+                }
+            },);
+
+            window.addEventListener("keyup", function (event) {
+                if (event.key === workflowHistory.prevKey){
+                    clearTimeout(workflowHistory.keyTimeout);
+                    workflowHistory.prevKey=null;
+                    workflowHistory.repKeyCount=0;
+                }
+            });
+
+            const o_onBeforeChange = app.graph.onBeforeChange;
+            app.graph.onBeforeChange = function (info) {
+                o_onBeforeChange ? o_onBeforeChange.apply(this, arguments) : undefined;
+
+                this.b_count += 1;
+                if (!workflowHistory.enabled) return;
+                if(this.b_count<2) return;//avoid unnecessary spam
+
+                //console.log("beforeChange");
+                workflowHistory.before();
+            }
+
+            const o_onAfterChange = app.graph.onAfterChange;
+            app.graph.afterChange = function (info) {
+                o_onAfterChange ? o_onAfterChange.apply(this, arguments) : undefined;
+
+                this.b_count -= 1;
+                if (!workflowHistory.enabled) return;
+                if(this.b_count>0) return; //same reasoning as in beforeChange
+
+                //console.log("afterChange");
+                workflowHistory.after();
+            }
+        }
+    }
+
+    before() {
+        if (!this.enabled) return;
+        //console.log("before");
+
+        const timestamp = Date.now();
+        if (timestamp - this.prev_undo_timestamp < this.state_merge_threshold) {
+            this.prev_undo_timestamp = timestamp;
+            //console.log("merged") // will be kept the same
+            return; //state is whithin the merge thesh and is discarded, don't do anything else
+        }
+
+        // potential state outside merge thresh. note that it is not guaranteed to be pushed        
+        this.get_new_candidate_state(app);
+    }
+
+    after() {
+        if (!this.enabled) return;
+        //console.log("after");
+
+        {
+            // was there any change? if not, don't try to store the state
+            // (this check is not guaranteed to work due to node order, but I am reordering nodes onSelected;
+            //  so unless there are other operations that reorder nodes, it should be fine... I think)
+            const equal = this.equal_states(JSON.stringify(app.graph.serialize(), null), this.temp_state);
+            if (equal) return;
+        }
+
+        this.tryAddToUndoHistory();
+    }
+
+    tryAddToUndoHistory() {
+        if (!this.enabled) return;
+
+        const potential_new_state = this.temp_state;
+        const timestamp = this.temp_timestamp;
+
+        if (this.undo_history.length > 0 && this.equal_states(potential_new_state, this.undo_history[0])) {
+            //console.log("discarded potential undo state")
+            return; //the new state is equal to prev stored state.
+        }
+
+        //console.log("added state to undo history")
+        this.undo_history.unshift(potential_new_state);
+        this.prev_undo_timestamp = timestamp;
+
+        if (this.undo_history.length > this.max_undo_steps)
+            this.undo_history.pop();
+
+        this.redo_history = []
+    }
+
+    undo(app) { this.undo_redo(app, false); }
+    redo(app) { this.undo_redo(app, true); }
+
+    undo_redo(app, redo) {
+        if (!this.enabled) return;
+
+        this.disable_load_reset = true;
+
+        let timeline = redo ? this.redo_history : this.undo_history;
+        let opposite_timeline = redo ? this.undo_history : this.redo_history;
+        let operation_name = redo ? "redo" : "undo";
+        let max_opposite_timeline_size = redo ? this.max_undo_steps : this.max_redo_steps;
+
+        if (timeline.length === 0) {
+            console.log("Can't " + operation_name) //maybe alert instead?
+            return;
+        }
+
+        const prev_state = timeline.shift();
+        const current_state = JSON.stringify(app.graph.serialize(), null);
+
+        opposite_timeline.unshift(current_state);
+        if (opposite_timeline.length > max_opposite_timeline_size)
+            opposite_timeline.pop();
+
+        app.loadGraphData(JSON.parse(prev_state));
+        app.graph.setDirtyCanvas(true);
+
+        this.get_new_candidate_state(app);
+
+        this.disable_load_reset = false;
+    }
+
+    equal_states(a, b) {
+        if (a.length === b.length) {
+            return a.localeCompare(b) === 0;
+        }
+        return false;
+    }
+
+    get_new_candidate_state(app) {
+        this.temp_timestamp = Date.now();
+        this.temp_state = JSON.stringify(app.graph.serialize(), null);
+    }
+
+    clean_history() {
+        this.undo_history = [];
+        this.redo_history = [];
+    }
+
+    time_to_next_operation_repeat(){
+        const max_delay = 500;
+        const min_delay = 30;
+        const repeats_till_min_delay = 12;
+        const rtmd_squared = repeats_till_min_delay*repeats_till_min_delay;
+
+        return max_delay - 
+            (max_delay - min_delay)/rtmd_squared 
+            * Math.min(this.repKeyCount*this.repKeyCount, rtmd_squared);
+    }
+
+}
+
+export const workflowHistory = new WorkflowHistory();