From 32c6d1215b20c343629aa9158fdfa71bfbef6c23 Mon Sep 17 00:00:00 2001 From: bennykok Date: Wed, 28 Feb 2024 14:28:28 -0800 Subject: [PATCH] feat(plugin): streaming file type support, webp and jepg, quality settings --- comfy-nodes/output_websocket_image.py | 6 ++++-- globals.py | 5 ++++- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/comfy-nodes/output_websocket_image.py b/comfy-nodes/output_websocket_image.py index 751118f..db19020 100644 --- a/comfy-nodes/output_websocket_image.py +++ b/comfy-nodes/output_websocket_image.py @@ -17,6 +17,8 @@ class ComfyDeployWebscoketImageOutput: {"multiline": False, "default": "output_id"}, ), "images": ("IMAGE", ), + "file_type": (["WEBP", "PNG", "JPEG"], ), + "quality": ("INT", {"default": 80, "min": 1, "max": 100, "step": 1}), }, "optional": { "client_id": ( @@ -36,7 +38,7 @@ class ComfyDeployWebscoketImageOutput: CATEGORY = "output" - def run(self, output_id, images, client_id): + def run(self, output_id, images, file_type, quality, client_id): prompt_server = PromptServer.instance loop = prompt_server.loop @@ -48,7 +50,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, ["PNG", image, None], client_id) + schedule_coroutine_blocking(send_image, [file_type, image, None, quality], client_id) print("Image sent") return {"ui": {}} diff --git a/globals.py b/globals.py index 6bde3bc..f0a63fb 100644 --- a/globals.py +++ b/globals.py @@ -15,6 +15,7 @@ async def send_image(image_data, sid=None): image_type = image_data[0] image = image_data[1] max_size = image_data[2] + quality = image_data[3] if max_size is not None: if hasattr(Image, 'Resampling'): resampling = Image.Resampling.BILINEAR @@ -27,11 +28,13 @@ async def send_image(image_data, sid=None): type_num = 1 elif image_type == "PNG": type_num = 2 + elif image_type == "WEBP": + type_num = 3 bytesIO = BytesIO() header = struct.pack(">I", type_num) bytesIO.write(header) - image.save(bytesIO, format=image_type, quality=95, compress_level=1) + 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)