fix(plugin): output_id is also included in the binary data back

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
@@ -38,6 +38,16 @@ class ComfyDeployWebscoketImageOutput:
 
     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
         loop = prompt_server.loop
@@ -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": {}}

globals.py

@@ -11,7 +11,14 @@ class BinaryEventTypes:
     PREVIEW_IMAGE = 1
     UNENCODED_PREVIEW_IMAGE = 2
     
-async def send_image(image_data, sid=None):
+max_output_id_length = 24
+
+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)