init

.gitattributes

@@ -0,0 +1 @@
+*.png filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Haoming
+
+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

@@ -0,0 +1,23 @@
+# ComfyUI Floodgate
+This is an Extension for [ComfyUI](https://github.com/comfyanonymous/ComfyUI), which allows you to control the logic flow with just one click!
+
+<p align="center"><img src="workflow.png"></p>
+
+<h3 align="center">Motivation</h3>
+<p align="center">While <b>Hires. Fix</b> <i>(or similar workflows)</i> can significantly improve the output,
+it also takes a longer time to process. Thus, many will simply run the base resolution until a good seed is found before enabling it. 
+However, for <code>ComfyUI</code> this means connecting and disconnecting multiple nodes every single time...</p>
+
+<h3 align="center">Solution</h3>
+<p align="center">Introducing, <b>Floodgate</b>! <br>
+Simply click one button to toggle between the logic flows. No more connecting and disconnecting multiple nodes! </p>
+
+<h3 align="center">How to Use</h3>
+<p align="center">Connect the <code>LATENT</code> output from a Sampler to the input of the Floodgate, 
+then connect each path to the desired remaining workflow. Under <code>Queue Prompt</code>, you can toggle between the path to take.
+The path not chosen will simply not be executed. Also, since <code>ComfyUI</code> caches the intermediate data, 
+opening the Floodgate will not require the precedent nodes to be processed again!</p>
+
+### Note / ToDo
+- Only **one** Floodgate can exist in a Workflow *as of now*
+- Only supports `LATENT` connections *as of now*

__init__.py

@@ -0,0 +1,96 @@
+from .floodgate import FloodGate
+from aiohttp import web
+import execution
+import server
+
+
+floodgate_open = False
+
+WEB_DIRECTORY = "js"
+NODE_CLASS_MAPPINGS = {"FloodGate": FloodGate}
+NODE_DISPLAY_NAME_MAPPINGS = {"FloodGate": "Flood Gate"}
+
+
+def find_gate(prompt:dict) -> str:
+    '''Find the Unique ID of the Floodgate Node'''
+    gate_ID = None
+
+    for k, v in prompt.items():
+        if v["class_type"] == "FloodGate":
+            if gate_ID is None:
+                gate_ID = k
+            else:
+                print('[Warning] Multiple Floodgates Detected! This will most likely raise errors!')
+
+    return gate_ID
+
+def block_gate(prompt:dict, gate_ID:str) -> dict:
+    '''"Bypass" the Nodes that should be Blocked'''
+    nodes_affected = []
+
+    for node, data in prompt.items():
+        for k, v in data["inputs"].items():
+            if not isinstance(v, list):
+                continue
+
+            if gate_ID in v:
+                if (not floodgate_open) and (v[1] == 1):
+                    nodes_affected.append(node)
+                    break
+
+                if (floodgate_open) and (v[1] == 0):
+                    nodes_affected.append(node)
+                    break
+
+    for key in nodes_affected:
+        del prompt[key]
+
+    if len(nodes_affected) > 0:
+        return recursive_block_gate(prompt, nodes_affected)
+    else:
+        return prompt
+
+def recursive_block_gate(prompt:dict, node_IDs:list) -> dict:
+    '''Block the subsequent nodes of which source has been blocked'''
+    to_delete = []
+
+    for node, data in prompt.items():
+        for k, v in data["inputs"].items():
+            # Connection is always a List
+            if not isinstance(v, list):
+                continue
+
+            if any(ID in v for ID in node_IDs):
+                to_delete.append(node)
+                break
+
+    for key in to_delete:
+        del prompt[key]
+
+    if len(to_delete) > 0:
+        return recursive_block_gate(prompt, to_delete)
+    else:
+        return prompt
+
+
+@server.PromptServer.instance.routes.get("/floodgate")
+async def floodgate_toggle(_):
+    '''Toggle the Floodgate Status'''
+
+    global floodgate_open
+    floodgate_open = not floodgate_open
+
+    return web.json_response({"status" : floodgate_open})
+
+
+original_validate = execution.validate_prompt
+
+def hijack_validate(prompt):
+
+    gate_ID = find_gate(prompt)
+    if gate_ID is None:
+        return original_validate(prompt)
+
+    return original_validate(block_gate(prompt, gate_ID))
+
+execution.validate_prompt = hijack_validate

floodgate.py

@@ -0,0 +1,14 @@
+class FloodGate:
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return { "required": { "latent": ("LATENT",) } }
+
+    RETURN_TYPES = ("LATENT", "LATENT")
+    RETURN_NAMES = ("CLOSE", "OPEN")
+
+    FUNCTION = "gate"
+    CATEGORY = "latent"
+
+    def gate(self, latent):
+        return (latent, latent)

js/flood_gate.js

@@ -0,0 +1,24 @@
+import { app } from "../../scripts/app.js";
+import { api } from "../../scripts/api.js";
+
+app.registerExtension({
+	name: "Comfy.FloodGate",
+	async setup() {
+		const floodButton = document.createElement("button");
+		floodButton.textContent = "Floodgate: Close";
+		floodButton.style.margin = '-4px 0px 6px 0px';
+
+		floodButton.addEventListener("click", async () => {
+			const resp = await api.fetchApi('/floodgate');
+			const stat = await resp.json();
+
+			if (stat['status'])
+				floodButton.textContent = "Floodgate: Open";
+			else
+				floodButton.textContent = "Floodgate: Close";
+		});
+
+		const queueButton = document.getElementById("queue-button");
+		queueButton.after(floodButton);
+	}
+});

workflow.png

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:adecf8e8dfcf4a2df1671095f1b947e1f1c0dc1b429fddc3e2b249e934e0328a
+size 400490