From e5e62a4a795b15195f6ce50df1feb3dd5a716e10 Mon Sep 17 00:00:00 2001
From: Marco Tundo <marco@bymar.co>
Date: Sat, 7 Sep 2024 19:06:09 -0400
Subject: [PATCH] Added filetype selector to SendImageWebSocket

---
 nodes.py | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/nodes.py b/nodes.py
index 39edb5d..68da235 100644
--- a/nodes.py
+++ b/nodes.py
@@ -50,19 +50,23 @@ class LoadMaskBase64:
         if img.dim() == 3:  # RGB(A) input, use red channel
             img = img[:, :, 0]
         return (img.unsqueeze(0),)
-
-
+    
 class SendImageWebSocket:
     @classmethod
     def INPUT_TYPES(s):
-        return {"required": {"images": ("IMAGE",)}}
+        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):
+    def send_images(self, images, format):
         results = []
         for tensor in images:
             array = 255.0 * tensor.cpu().numpy()
@@ -71,17 +75,15 @@ class SendImageWebSocket:
             server = PromptServer.instance
             server.send_sync(
                 BinaryEventTypes.UNENCODED_PREVIEW_IMAGE,
-                ["PNG", image, None],
+                [format, image, None],
                 server.client_id,
             )
             results.append(
-                # Could put some kind of ID here, but for now just match them by index
-                {"source": "websocket", "content-type": "image/png", "type": "output"}
+                {"source": "websocket", "content-type": f"image/{format.lower()}", "type": "output"}
             )
 
         return {"ui": {"images": results}}
 
-
 class CropImage:
     """Deprecated, ComfyUI has an ImageCrop node now which does the same."""