feat(plugin): streaming file type support, webp and jepg, quality settings

This commit is contained in:
bennykok 2024-02-28 14:28:28 -08:00
parent 9e79c434a9
commit 32c6d1215b
2 changed files with 8 additions and 3 deletions

View File

@ -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": {}}

View File

@ -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)