From 410d03cd2bdd8daaa5e8ac498364cc5b20b98228 Mon Sep 17 00:00:00 2001
From: bennykok <itechbenny@gmail.com>
Date: Thu, 29 Feb 2024 11:39:54 -0800
Subject: [PATCH] fix(plugin): output_id is also included in the binary data
 back

---
 comfy-nodes/output_websocket_image.py | 14 ++++++++++++--
 globals.py                            | 17 ++++++++++++++++-
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/comfy-nodes/output_websocket_image.py b/comfy-nodes/output_websocket_image.py
index db19020..7c6b47a 100644
--- a/comfy-nodes/output_websocket_image.py
+++ b/comfy-nodes/output_websocket_image.py
@@ -5,7 +5,7 @@ import torch
 from server import PromptServer, BinaryEventTypes
 import asyncio
 
-from globals import send_image
+from globals import send_image, max_output_id_length
 
 class ComfyDeployWebscoketImageOutput:
     @classmethod
@@ -37,6 +37,16 @@ class ComfyDeployWebscoketImageOutput:
     FUNCTION = "run"
 
     CATEGORY = "output"
+    
+    @classmethod
+    def VALIDATE_INPUTS(s, output_id):
+        try:
+            if len(output_id.encode('ascii')) > max_output_id_length:
+                raise ValueError(f"output_id size is greater than {max_output_id_length} bytes")
+        except UnicodeEncodeError:
+            raise ValueError("output_id is not ASCII encodable")
+
+        return True
 
     def run(self, output_id, images, file_type, quality, client_id):
         prompt_server = PromptServer.instance
@@ -50,7 +60,7 @@ class ComfyDeployWebscoketImageOutput:
             array = 255.0 * tensor.cpu().numpy()
             image = Image.fromarray(np.clip(array, 0, 255).astype(np.uint8))
 
-            schedule_coroutine_blocking(send_image, [file_type, image, None, quality], client_id)
+            schedule_coroutine_blocking(send_image, [file_type, image, None, quality], client_id, output_id)
             print("Image sent")
 
         return {"ui": {}}
diff --git a/globals.py b/globals.py
index f0a63fb..faba46c 100644
--- a/globals.py
+++ b/globals.py
@@ -10,8 +10,15 @@ sockets = dict()
 class BinaryEventTypes:
     PREVIEW_IMAGE = 1
     UNENCODED_PREVIEW_IMAGE = 2
+    
+max_output_id_length = 24
 
-async def send_image(image_data, sid=None):
+async def send_image(image_data, sid=None, output_id:str = None):
+    max_length = max_output_id_length
+    output_id = output_id[:max_length]
+    padded_output_id = output_id.ljust(max_length, '\x00')
+    encoded_output_id = padded_output_id.encode('ascii', 'replace')
+    
     image_type = image_data[0]
     image = image_data[1]
     max_size = image_data[2]
@@ -33,7 +40,15 @@ async def send_image(image_data, sid=None):
 
     bytesIO = BytesIO()
     header = struct.pack(">I", type_num)
+    # 4 bytes for the type
     bytesIO.write(header)
+    # 10 bytes for the output_id
+    position_before = bytesIO.tell()
+    bytesIO.write(encoded_output_id)
+    position_after = bytesIO.tell()
+    bytes_written = position_after - position_before
+    print(f"Bytes written: {bytes_written}")
+    
     image.save(bytesIO, format=image_type, quality=quality, compress_level=1)
     preview_bytes = bytesIO.getvalue()
     await send_bytes(BinaryEventTypes.PREVIEW_IMAGE, preview_bytes, sid=sid)