diff --git a/README.md b/README.md index a4f4e0d..e653b50 100644 --- a/README.md +++ b/README.md @@ -1,23 +1,28 @@ # 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! +This is an Extension for [ComfyUI](https://github.com/comfyanonymous/ComfyUI), which allows you to easily control the logic flow! -

+## New Features +1. You can now connect any arbitrary types + - *(input/output types still have to match)* +2. You can now have multiple Floodgates in one workflow +3. Control each Floodgate individually on itself + +

+ +

Motivation

-### Motivation While **Hires. Fix** *(or similar workflows)* 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 `ComfyUI` this means connecting and disconnecting multiple nodes every single time... -### Solution +

Solution

+ Introducing, **Floodgate**!
-Simply click one button to toggle between the logic flows. No more connecting and disconnecting multiple nodes! +Simply toggle between the logic flows. No more reconnecting multiple nodes! -### How to Use -Connect the `LATENT` output from a Sampler node to the input of the Floodgate, -then connect each path to the desired remaining workflow. Under `Queue Prompt`, you can toggle between the path to take. -The path not chosen will simply not be executed. Also, since `ComfyUI` caches the intermediate data, +

How to Use

+ +Connect the output of a node *(**eg.** `LATENT`)* to the `source` of the Floodgate node, +then connect each path to the desired remaining workflow. Simply click on `gate_open` to toggle between the paths to take. +The path not chosen will not be executed. Furthermore, since `ComfyUI` caches the intermediate results, opening the Floodgate will not require the precedent nodes to be processed again! - -### Note / ToDo -- Only **one** Floodgate can exist in a Workflow *as of now* -- Only supports `LATENT` connections *as of now* diff --git a/__init__.py b/__init__.py index 49e5cd6..b91538d 100644 --- a/__init__.py +++ b/__init__.py @@ -2,38 +2,51 @@ from .floodgate import FloodGate from aiohttp import web import execution import server +import nodes - -floodgate_open = False - -WEB_DIRECTORY = "js" NODE_CLASS_MAPPINGS = {"FloodGate": FloodGate} NODE_DISPLAY_NAME_MAPPINGS = {"FloodGate": "Flood Gate"} -def find_gate(prompt:dict) -> str: +def find_gate(prompt:dict) -> list: '''Find the Unique ID of the Floodgate Node''' - gate_ID = None + gate_IDs = [] 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!') + gate_IDs.append(k) - return gate_ID + # if len(gate_IDs) > 1: + # print('[Warning] Multiple Floodgates Detected is still experimental!') -def block_gate(prompt:dict, gate_ID:str) -> dict: + return gate_IDs + +def block_gate(prompt:dict, gate_ID:str, floodgate_open:bool) -> dict: '''"Bypass" the Nodes that should be Blocked''' nodes_affected = [] + try: + sauce_id, out_index = prompt[gate_ID]['inputs']['source'] + except KeyError: + # Floodgate is not connected; let ComfyUI raise the error + return prompt + + sauce_class = nodes.NODE_CLASS_MAPPINGS[prompt[sauce_id]['class_type']] + sauce_type = str(sauce_class.RETURN_TYPES[out_index]).lower().strip() + for node, data in prompt.items(): for k, v in data["inputs"].items(): if not isinstance(v, list): continue if gate_ID in v: + + target_class = nodes.NODE_CLASS_MAPPINGS[data['class_type']] + target_type = (target_class.INPUT_TYPES()['required'][k][0]).lower().strip() + + if sauce_type != target_type: + raise IOError() + if (not floodgate_open) and (v[1] == 1): nodes_affected.append(node) break @@ -71,26 +84,34 @@ def recursive_block_gate(prompt:dict, node_IDs:list) -> dict: 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_IDs:list = find_gate(prompt) - gate_ID = find_gate(prompt) - if gate_ID is None: + if len(gate_IDs) == 0: return original_validate(prompt) - return original_validate(block_gate(prompt, gate_ID)) + for ID in gate_IDs: + if ID not in prompt.keys(): + continue + + try: + gate_open = prompt[ID]['inputs']['gate_open'] + prompt = block_gate(prompt, ID, gate_open) + + except IOError: + return (False, + { + 'type': 'floodgate_io_mismatch', + 'message': 'Floodgate IO Type Mismatch', + 'details': 'source cannot be connected to outputs', + 'extra_info': {} + } + , [], []) + + return original_validate(prompt) execution.validate_prompt = hijack_validate diff --git a/example.png b/example.png deleted file mode 100644 index fb7d5b4..0000000 Binary files a/example.png and /dev/null differ diff --git a/floodgate.py b/floodgate.py index 199dc87..93abaf6 100644 --- a/floodgate.py +++ b/floodgate.py @@ -1,14 +1,31 @@ -class FloodGate: +class AnyType(str): + """A special class that is always equal in comparisons. Credit to crystian & pythongosssss""" + def __eq__(self, _) -> bool: + return True + + def __ne__(self, __value: object) -> bool: + return False + + +generic = AnyType("*") + + +class FloodGate: @classmethod def INPUT_TYPES(s): - return { "required": { "latent": ("LATENT",) } } + return { + "required": { + "source": (generic,), + "gate_open": ("BOOLEAN", {"default": False}), + } + } - RETURN_TYPES = ("LATENT", "LATENT") + RETURN_TYPES = (generic, generic) RETURN_NAMES = ("CLOSE", "OPEN") FUNCTION = "gate" CATEGORY = "latent" - def gate(self, latent): - return (latent, latent) + def gate(self, source, gate_open): + return (source, source) diff --git a/js/flood_gate.js b/js/flood_gate.js deleted file mode 100644 index 57f74f0..0000000 --- a/js/flood_gate.js +++ /dev/null @@ -1,24 +0,0 @@ -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); - } -}); diff --git a/workflow.png b/workflow.png new file mode 100644 index 0000000..83722d7 Binary files /dev/null and b/workflow.png differ