From 212d6b7728fca6143dd7c4e62dbb11107ef75995 Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Thu, 13 Aug 2026 05:39:24 +0200 Subject: [PATCH] refactor: share backend image serialization helpers --- python/image_serialization.py | 8 +- python/image_utils.py | 9 +-- python/node.py | 88 +++++++++++----------- tests/test_entrypoint_integration.py | 108 +++++++++++++++++++++++++++ tests/test_image_serialization.py | 12 ++- 5 files changed, 170 insertions(+), 55 deletions(-) diff --git a/python/image_serialization.py b/python/image_serialization.py index ce8c08a..c306b68 100644 --- a/python/image_serialization.py +++ b/python/image_serialization.py @@ -29,4 +29,10 @@ def file_to_data_url(path: str, *, mime_type: str = "image/png") -> str: return _bytes_to_data_url(image_file.read(), mime_type) -__all__ = ["file_to_data_url", "pil_to_data_url"] +def data_url_to_pil(data_url: str) -> Image.Image: + """Decode an image data URL into a lazily opened PIL image.""" + image_bytes = base64.b64decode(data_url.split(",")[1]) + return Image.open(io.BytesIO(image_bytes)) + + +__all__ = ["data_url_to_pil", "file_to_data_url", "pil_to_data_url"] diff --git a/python/image_utils.py b/python/image_utils.py index ca818cd..c4fe6d9 100644 --- a/python/image_utils.py +++ b/python/image_utils.py @@ -1,22 +1,17 @@ """Image and tensor conversion helpers shared by LayerForge endpoints.""" -import base64 -import io - import numpy as np from PIL import Image - -from .image_serialization import pil_to_data_url from torchvision import transforms +from .image_serialization import data_url_to_pil, pil_to_data_url from .node import log def convert_base64_to_tensor(base64_str): """Convert a data-URL image into a BCHW tensor and optional alpha tensor.""" try: - img_data = base64.b64decode(base64_str.split(",")[1]) - img = Image.open(io.BytesIO(img_data)) + img = data_url_to_pil(base64_str) has_alpha = img.mode == "RGBA" alpha = None diff --git a/python/node.py b/python/node.py index b6b607d..baaadbf 100644 --- a/python/node.py +++ b/python/node.py @@ -1,7 +1,5 @@ """Core LayerForge ComfyUI node and its in-memory canvas state.""" -import base64 -import io import os import threading import time @@ -15,7 +13,9 @@ import torch import torch.nn.functional as F from PIL import Image -from .image_serialization import pil_to_data_url +from .image_serialization import data_url_to_pil, pil_to_data_url + +_OUTPUT_IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg", ".bmp", ".gif") try: from .log_system import create_module_logger @@ -172,6 +172,16 @@ class LayerForgeNode: FUNCTION = "process_canvas_image" CATEGORY = "azNodes > LayerForge" + @staticmethod + def _serialize_rgb_tensor_sample(image_tensor): + image_array = (image_tensor.cpu().numpy() * 255).astype(np.uint8) + pil_image = Image.fromarray(image_array, "RGB") + return { + "data": pil_to_data_url(pil_image), + "width": pil_image.width, + "height": pil_image.height, + } + def add_image_to_canvas(self, input_image): try: if not isinstance(input_image, torch.Tensor): @@ -252,27 +262,21 @@ class LayerForgeNode: log.info(f"Processing batch of {batch_size} image(s)") if batch_size == 1: - img_np = (input_image.squeeze(0).cpu().numpy() * 255).astype(np.uint8) - pil_img = Image.fromarray(img_np, "RGB") - input_data["input_image"] = pil_to_data_url(pil_img) - input_data["input_image_width"] = pil_img.width - input_data["input_image_height"] = pil_img.height - log.debug(f"Stored single input image: {pil_img.width}x{pil_img.height}") + serialized_image = self._serialize_rgb_tensor_sample(input_image.squeeze(0)) + input_data["input_image"] = serialized_image["data"] + input_data["input_image_width"] = serialized_image["width"] + input_data["input_image_height"] = serialized_image["height"] + log.debug( + f"Stored single input image: {serialized_image['width']}x{serialized_image['height']}" + ) else: images_array = [] for index in range(batch_size): - img_np = (input_image[index].cpu().numpy() * 255).astype(np.uint8) - pil_img = Image.fromarray(img_np, "RGB") - images_array.append( - { - "data": pil_to_data_url(pil_img), - "width": pil_img.width, - "height": pil_img.height, - } - ) + serialized_image = self._serialize_rgb_tensor_sample(input_image[index]) + images_array.append(serialized_image) log.debug( f"Stored batch image {index + 1}/{batch_size}: " - f"{pil_img.width}x{pil_img.height}" + f"{serialized_image['width']}x{serialized_image['height']}" ) input_data["input_images_batch"] = images_array @@ -300,15 +304,13 @@ class LayerForgeNode: if canvas_data: log.info(f"Canvas data found for node {node_id} from WebSocket") if canvas_data.get("image"): - image_bytes = base64.b64decode(canvas_data["image"].split(",")[1]) - pil_image = Image.open(io.BytesIO(image_bytes)).convert("RGB") + pil_image = data_url_to_pil(canvas_data["image"]).convert("RGB") image_array = np.array(pil_image).astype(np.float32) / 255.0 processed_image = torch.from_numpy(image_array)[None,] log.debug(f"Image loaded from WebSocket, shape: {processed_image.shape}") if canvas_data.get("mask"): - mask_bytes = base64.b64decode(canvas_data["mask"].split(",")[1]) - pil_mask = Image.open(io.BytesIO(mask_bytes)).convert("L") + pil_mask = data_url_to_pil(canvas_data["mask"]).convert("L") mask_array = np.array(pil_mask).astype(np.float32) / 255.0 processed_mask = torch.from_numpy(mask_array)[None,] log.debug(f"Mask loaded from WebSocket, shape: {processed_mask.shape}") @@ -353,38 +355,33 @@ class LayerForgeNode: @classmethod def get_latest_image(cls): - output_dir = folder_paths.get_output_directory() - files = [ - os.path.join(output_dir, filename) - for filename in os.listdir(output_dir) - if os.path.isfile(os.path.join(output_dir, filename)) - ] - image_files = [ - path for path in files if path.lower().endswith((".png", ".jpg", ".jpeg", ".bmp", ".gif")) - ] + image_files = list(cls._iter_output_image_paths()) if not image_files: return None return max(image_files, key=os.path.getctime) @classmethod def get_latest_images(cls, since_timestamp=0): - output_dir = folder_paths.get_output_directory() files = [] - for filename in os.listdir(output_dir): - file_path = os.path.join(output_dir, filename) - if os.path.isfile(file_path) and file_path.lower().endswith( - (".png", ".jpg", ".jpeg", ".bmp", ".gif") - ): - try: - mtime = os.path.getmtime(file_path) - if mtime > since_timestamp: - files.append((mtime, file_path)) - except OSError: - continue + for file_path in cls._iter_output_image_paths(): + try: + mtime = os.path.getmtime(file_path) + if mtime > since_timestamp: + files.append((mtime, file_path)) + except OSError: + continue files.sort(key=lambda item: item[0]) return [item[1] for item in files] + @classmethod + def _iter_output_image_paths(cls): + output_dir = folder_paths.get_output_directory() + for filename in os.listdir(output_dir): + file_path = os.path.join(output_dir, filename) + if os.path.isfile(file_path) and file_path.lower().endswith(_OUTPUT_IMAGE_EXTENSIONS): + yield file_path + @classmethod def get_flow_status(cls, flow_id=None): if flow_id: @@ -423,8 +420,7 @@ class LayerForgeNode: def store_image(self, image_data): if isinstance(image_data, str) and image_data.startswith("data:image"): - image_bytes = base64.b64decode(image_data.split(",")[1]) - self.cached_image = Image.open(io.BytesIO(image_bytes)) + self.cached_image = data_url_to_pil(image_data) else: self.cached_image = image_data diff --git a/tests/test_entrypoint_integration.py b/tests/test_entrypoint_integration.py index 88c1ce3..2d9af75 100644 --- a/tests/test_entrypoint_integration.py +++ b/tests/test_entrypoint_integration.py @@ -203,6 +203,114 @@ def test_base64_image_conversion_preserves_rgb_and_alpha(layerforge_runtime): assert roundtrip_image.getpixel((0, 0)) == (255, 0, 0) +def test_tensor_input_payloads_preserve_single_and_batch_image_contract(layerforge_runtime): + import torch + from PIL import Image + + node_class = layerforge_runtime.node.LayerForgeNode + node_class._canvas_data_storage.clear() + node_class._canvas_cache["persistent_cache"] = {} + node = node_class() + single_image = torch.tensor( + [[ + [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + [[0.0, 0.0, 1.0], [1.0, 1.0, 1.0]], + ]], + dtype=torch.float32, + ) + + node.process_canvas_image(False, False, False, 0, "single-input", input_image=single_image) + single_payload = node_class._canvas_data_storage["single-input_input"] + single_image_data = base64.b64decode(single_payload["input_image"].split(",", 1)[1]) + single_decoded = Image.open(io.BytesIO(single_image_data)) + + assert single_payload["input_image_width"] == 2 + assert single_payload["input_image_height"] == 2 + assert single_decoded.size == (2, 2) + assert single_decoded.getpixel((0, 0)) == (255, 0, 0) + + batch_image = torch.stack((single_image[0], single_image[0] * 0.5)) + node.process_canvas_image(False, False, False, 0, "batch-input", input_image=batch_image) + batch_payload = node_class._canvas_data_storage["batch-input_input"] + + assert "input_image" not in batch_payload + assert [item["width"] for item in batch_payload["input_images_batch"]] == [2, 2] + assert [item["height"] for item in batch_payload["input_images_batch"]] == [2, 2] + assert all(item["data"].startswith("data:image/png;base64,") for item in batch_payload["input_images_batch"]) + + +def test_latest_image_helpers_preserve_filtering_and_time_semantics(layerforge_runtime, monkeypatch, tmp_path): + node_module = layerforge_runtime.node + node_class = node_module.LayerForgeNode + output_dir = tmp_path / "output" + output_dir.mkdir() + first_image = output_dir / "first.png" + second_image = output_dir / "second.JPG" + ignored_file = output_dir / "ignored.txt" + ignored_directory = output_dir / "folder.png" + first_image.write_bytes(b"first") + second_image.write_bytes(b"second") + ignored_file.write_bytes(b"ignored") + ignored_directory.mkdir() + + monkeypatch.setattr( + node_module.folder_paths, + "get_output_directory", + lambda: str(output_dir), + raising=False, + ) + creation_times = { + str(first_image): 10, + str(second_image): 20, + } + monkeypatch.setattr(node_module.os.path, "getctime", lambda path: creation_times[path]) + assert node_class.get_latest_image() == str(second_image) + + modification_times = { + str(first_image): 30, + str(second_image): 20, + } + monkeypatch.setattr(node_module.os.path, "getmtime", lambda path: modification_times[path]) + assert node_class.get_latest_images(0) == [str(second_image), str(first_image)] + assert node_class.get_latest_images(25) == [str(first_image)] + + +def test_websocket_and_cached_image_decoding_preserve_image_and_mask_modes(layerforge_runtime): + from PIL import Image + + node_class = layerforge_runtime.node.LayerForgeNode + node_class._canvas_data_storage.clear() + node_class._canvas_cache["persistent_cache"] = {} + + image = Image.new("RGBA", (2, 1), (255, 0, 0, 128)) + mask = Image.new("L", (2, 1), 128) + + def encode(pil_image): + buffer = io.BytesIO() + pil_image.save(buffer, format="PNG") + return "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode() + + node_class._canvas_data_storage["decode-node"] = { + "image": encode(image), + "mask": encode(mask), + } + node = node_class() + processed_image, processed_mask = node.process_canvas_image( + False, + False, + False, + 0, + "decode-node", + ) + + assert tuple(processed_image.shape) == (1, 1, 2, 3) + assert tuple(processed_mask.shape) == (1, 1, 2) + + node.store_image(encode(image)) + assert node.cached_image.mode == "RGBA" + assert node.cached_image.size == (2, 1) + + def test_matting_adapter_uses_native_loader_and_bhwc_input(layerforge_runtime, monkeypatch): import torch diff --git a/tests/test_image_serialization.py b/tests/test_image_serialization.py index afd92da..0f732d3 100644 --- a/tests/test_image_serialization.py +++ b/tests/test_image_serialization.py @@ -3,7 +3,7 @@ import io from PIL import Image -from python.image_serialization import file_to_data_url, pil_to_data_url +from python.image_serialization import data_url_to_pil, file_to_data_url, pil_to_data_url def _decode_data_url(data_url): @@ -33,3 +33,13 @@ def test_file_to_data_url_preserves_file_bytes_and_mime_type(tmp_path): assert header == "data:image/png;base64" assert encoded == b"not-transformed-image-bytes" + + +def test_data_url_to_pil_preserves_image_mode_dimensions_and_pixels(): + image = Image.new("RGBA", (2, 1), (255, 0, 128, 128)) + + decoded = data_url_to_pil(pil_to_data_url(image)) + + assert decoded.mode == "RGBA" + assert decoded.size == (2, 1) + assert decoded.getpixel((0, 0)) == (255, 0, 128, 128)