Multi-Process Parallelization with ProcessPoolExecutor

nodes/__init__.py

@@ -3,6 +3,7 @@
 from .audio_utils import *
 from .tensor_utils import *
 from .video_stream_utils import *
+from .native_utils import *
 from .api import *
 from .web import *
 
@@ -11,7 +12,7 @@
 NODE_DISPLAY_NAME_MAPPINGS = {}
 
 # Import and update mappings from submodules
-for module in [audio_utils, tensor_utils, video_stream_utils, api, web]:
+for module in [audio_utils, tensor_utils, video_stream_utils, api, web, native_utils]:
     if hasattr(module, 'NODE_CLASS_MAPPINGS'):
         NODE_CLASS_MAPPINGS.update(module.NODE_CLASS_MAPPINGS)
     if hasattr(module, 'NODE_DISPLAY_NAME_MAPPINGS'):

nodes/native_utils/__init__.py

@@ -0,0 +1,20 @@
+from .load_image_base64 import LoadImageBase64
+from .send_image_websocket import SendImageWebsocket
+from .send_tensor_websocket import SendTensorWebSocket
+
+# This dictionary is used by ComfyUI to register the nodes
+NODE_CLASS_MAPPINGS = {
+    "LoadImageBase64": LoadImageBase64, 
+    "SendImageWebsocket": SendImageWebsocket,
+    "SendTensorWebSocket": SendTensorWebSocket
+}
+
+# This dictionary provides display names for the nodes in the UI
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "LoadImageBase64": "Load Image Base64 (ComfyStream)",
+    "SendImageWebsocket": "Send Image Websocket (ComfyStream)",
+    "SendTensorWebSocket": "Save Tensor WebSocket (ComfyStream)"
+}
+
+# Export these variables for ComfyUI to use
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

nodes/native_utils/load_image_base64.py

@@ -0,0 +1,37 @@
+# borrowed from Acly's comfyui-tooling-nodes
+# https://github.com/Acly/comfyui-tooling-nodes/blob/main/nodes.py
+
+# TODO: I think we can recieve tensor data directly from the pipeline through the /prompt endpoint as JSON
+#       This may be more efficient than sending base64 encoded images through the websocket and
+#       allow for alternative data formats.
+
+from PIL import Image
+import base64
+import numpy as np
+import torch
+from io import BytesIO
+
+class LoadImageBase64:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"image": ("STRING", {"multiline": False})}}
+
+    RETURN_TYPES = ("IMAGE", "MASK")
+    CATEGORY = "external_tooling"
+    FUNCTION = "load_image"
+
+    def load_image(self, image):
+        imgdata = base64.b64decode(image)
+        img = Image.open(BytesIO(imgdata))
+
+        if "A" in img.getbands():
+            mask = np.array(img.getchannel("A")).astype(np.float32) / 255.0
+            mask = torch.from_numpy(mask)
+        else:
+            mask = None
+
+        img = img.convert("RGB")
+        img = np.array(img).astype(np.float32) / 255.0
+        img = torch.from_numpy(img)[None,]
+
+        return (img, mask)

nodes/native_utils/send_image_websocket.py

@@ -0,0 +1,44 @@
+# borrowed from Acly's comfyui-tooling-nodes
+# https://github.com/Acly/comfyui-tooling-nodes/blob/main/nodes.py
+
+# TODO: I think we can send tensor data directly to the pipeline in the websocket response.
+#       Worth talking to ComfyAnonymous about this.
+
+import numpy as np
+from PIL import Image
+from server import PromptServer, BinaryEventTypes
+
+class SendImageWebsocket:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "format": (["PNG", "JPEG"], {"default": "PNG"}),
+            }
+        }
+
+    RETURN_TYPES = ()
+    FUNCTION = "send_images"
+    OUTPUT_NODE = True
+    CATEGORY = "external_tooling"
+
+    def send_images(self, images, format):
+        results = []
+        for tensor in images:
+            array = 255.0 * tensor.cpu().numpy()
+            image = Image.fromarray(np.clip(array, 0, 255).astype(np.uint8))
+
+            server = PromptServer.instance
+            server.send_sync(
+                BinaryEventTypes.UNENCODED_PREVIEW_IMAGE,
+                [format, image, None],
+                server.client_id,
+            )
+            results.append({
+                "source": "websocket",
+                "content-type": f"image/{format.lower()}",
+                "type": "output",
+            })
+
+        return {"ui": {"images": results}}