From 32034217ae14bb0dac3108b1e4e8454f036a9b2b Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Sat, 13 Jan 2024 18:57:47 +0000
Subject: [PATCH] add setting to change control after generate to run before

---
 web/scripts/app.js     | 11 +++++++-
 web/scripts/widgets.js | 64 +++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 73 insertions(+), 2 deletions(-)

diff --git a/web/scripts/app.js b/web/scripts/app.js
index 72f9e860..d131045d 100644
--- a/web/scripts/app.js
+++ b/web/scripts/app.js
@@ -1,5 +1,5 @@
 import { ComfyLogging } from "./logging.js";
-import { ComfyWidgets } from "./widgets.js";
+import { ComfyWidgets, initWidgets } from "./widgets.js";
 import { ComfyUI, $el } from "./ui.js";
 import { api } from "./api.js";
 import { defaultGraph } from "./defaultGraph.js";
@@ -1420,6 +1420,7 @@ export class ComfyApp {
 
 		await this.#invokeExtensionsAsync("init");
 		await this.registerNodes();
+		initWidgets(this);
 
 		// Load previous workflow
 		let restored = false;
@@ -1774,6 +1775,14 @@ export class ComfyApp {
 	 */
 	async graphToPrompt() {
 		for (const outerNode of this.graph.computeExecutionOrder(false)) {
+			if (outerNode.widgets) {
+				for (const widget of outerNode.widgets) {
+					// Allow widgets to run callbacks before a prompt has been queued
+					// e.g. random seed before every gen
+					widget.beforeQueued?.();
+				}
+			}
+
 			const innerNodes = outerNode.getInnerNodes ? outerNode.getInnerNodes() : [outerNode];
 			for (const node of innerNodes) {
 				if (node.isVirtualNode) {
diff --git a/web/scripts/widgets.js b/web/scripts/widgets.js
index e2e21164..36429266 100644
--- a/web/scripts/widgets.js
+++ b/web/scripts/widgets.js
@@ -1,6 +1,19 @@
 import { api } from "./api.js"
 import "./domWidget.js";
 
+let controlValueRunBefore = false;
+function updateControlWidgetLabel(widget) {
+	let replacement = "after";
+	let find = "before";
+	if (controlValueRunBefore) {
+		[find, replacement] = [replacement, find]
+	}
+	widget.label = (widget.label ?? widget.name).replace(find, replacement);
+}
+
+const IS_CONTROL_WIDGET = Symbol();
+const HAS_EXECUTED = Symbol();
+
 function getNumberDefaults(inputData, defaultStep, precision, enable_rounding) {
 	let defaultVal = inputData[1]["default"];
 	let { min, max, step, round} = inputData[1];
@@ -62,6 +75,8 @@ export function addValueControlWidgets(node, targetWidget, defaultValue = "rando
 			serialize: false, // Don't include this in prompt.
 		}
 	);
+	valueControl[IS_CONTROL_WIDGET] = true;
+	updateControlWidgetLabel(valueControl);
 	widgets.push(valueControl);
 
 	const isCombo = targetWidget.type === "combo";
@@ -76,10 +91,12 @@ export function addValueControlWidgets(node, targetWidget, defaultValue = "rando
 				serialize: false, // Don't include this in prompt.
 			}
 		);
+		updateControlWidgetLabel(comboFilter);
+
 		widgets.push(comboFilter);
 	}
 
-	valueControl.afterQueued = () => {
+	const applyWidgetControl = () => {
 		var v = valueControl.value;
 
 		if (isCombo && v !== "fixed") {
@@ -159,6 +176,23 @@ export function addValueControlWidgets(node, targetWidget, defaultValue = "rando
 			targetWidget.callback(targetWidget.value);
 		}
 	};
+
+	valueControl.beforeQueued = () => {
+		if (controlValueRunBefore) {
+			// Don't run on first execution
+			if (valueControl[HAS_EXECUTED]) {
+				applyWidgetControl();
+			}
+		}
+		valueControl[HAS_EXECUTED] = true;
+	};
+
+	valueControl.afterQueued = () => {
+		if (!controlValueRunBefore) {
+			applyWidgetControl();
+		}
+	};
+
 	return widgets;
 };
 
@@ -224,6 +258,34 @@ function isSlider(display, app) {
 	return (display==="slider") ? "slider" : "number"
 }
 
+export function initWidgets(app) {
+	app.ui.settings.addSetting({
+		id: "Comfy.WidgetControlMode",
+		name: "Widget Value Control Mode",
+		type: "combo",
+		defaultValue: "after",
+		options: ["before", "after"],
+		tooltip: "Controls when widget values are updated (randomize/increment/decrement), either before the prompt is queued or after.",
+		onChange(value) {
+			controlValueRunBefore = value === "before";
+			for (const n of app.graph._nodes) {
+				if (!n.widgets) continue;
+				for (const w of n.widgets) {
+					if (w[IS_CONTROL_WIDGET]) {
+						updateControlWidgetLabel(w);
+						if (w.linkedWidgets) {
+							for (const l of w.linkedWidgets) {
+								updateControlWidgetLabel(l);
+							}
+						}
+					}
+				}
+			}
+			app.graph.setDirtyCanvas(true);
+		},
+	});
+}
+
 export const ComfyWidgets = {
 	"INT:seed": seedWidget,
 	"INT:noise_seed": seedWidget,