Major Rewrite

addressing #1

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!
 
-<p align="center"><img src="example.png"></p>
+## 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
+
+<p align="center"><img src="workflow.png"></p>
+
+<h3 align="center">Motivation</h3>
 
-### 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
+<h3 align="center">Solution</h3>
+
 Introducing, **Floodgate**! <br>
-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, 
+<h3 align="center">How to Use</h3>
+
+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*

__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

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)

js/flood_gate.js

@@ -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);
-	}
-});