refactor: centralize image loading and serialization

python/image_serialization.py

@@ -0,0 +1,32 @@
+"""Shared image serialization helpers for LayerForge backend boundaries."""
+
+import base64
+import io
+
+from PIL import Image
+
+
+def _bytes_to_data_url(data: bytes, mime_type: str) -> str:
+    encoded = base64.b64encode(data).decode("ascii")
+    return f"data:{mime_type};base64,{encoded}"
+
+
+def pil_to_data_url(
+    image: Image.Image,
+    *,
+    image_format: str = "PNG",
+    mime_type: str = "image/png",
+) -> str:
+    """Serialize a PIL image using the existing PNG data-URL contract."""
+    buffer = io.BytesIO()
+    image.save(buffer, format=image_format)
+    return _bytes_to_data_url(buffer.getvalue(), mime_type)
+
+
+def file_to_data_url(path: str, *, mime_type: str = "image/png") -> str:
+    """Read an image file without transforming its bytes into a data URL."""
+    with open(path, "rb") as image_file:
+        return _bytes_to_data_url(image_file.read(), mime_type)
+
+
+__all__ = ["file_to_data_url", "pil_to_data_url"]

python/image_utils.py

@@ -5,6 +5,8 @@ import io
 
 import numpy as np
 from PIL import Image
+
+from .image_serialization import pil_to_data_url
 from torchvision import transforms
 
 from .node import log
@@ -61,10 +63,7 @@ def convert_tensor_to_base64(tensor, alpha_mask=None, original_alpha=None):
         else:
             img = Image.fromarray(img_array, mode="RGB")
 
-        buffer = io.BytesIO()
-        img.save(buffer, format="PNG")
-        img_str = base64.b64encode(buffer.getvalue()).decode()
-        return f"data:image/png;base64,{img_str}"
+        return pil_to_data_url(img)
     except Exception as error:
         log.error(f"Error in convert_tensor_to_base64: {error}")
         log.debug(f"Tensor shape: {tensor.shape}, dtype: {tensor.dtype}")

python/node.py

@@ -15,6 +15,8 @@ import torch
 import torch.nn.functional as F
 from PIL import Image
 
+from .image_serialization import pil_to_data_url
+
 try:
     from .log_system import create_module_logger
 
@@ -252,10 +254,7 @@ class LayerForgeNode:
                     if batch_size == 1:
                         img_np = (input_image.squeeze(0).cpu().numpy() * 255).astype(np.uint8)
                         pil_img = Image.fromarray(img_np, "RGB")
-                        buffered = io.BytesIO()
-                        pil_img.save(buffered, format="PNG")
-                        img_str = base64.b64encode(buffered.getvalue()).decode()
-                        input_data["input_image"] = f"data:image/png;base64,{img_str}"
+                        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}")
@@ -264,12 +263,9 @@ class LayerForgeNode:
                         for index in range(batch_size):
                             img_np = (input_image[index].cpu().numpy() * 255).astype(np.uint8)
                             pil_img = Image.fromarray(img_np, "RGB")
-                            buffered = io.BytesIO()
-                            pil_img.save(buffered, format="PNG")
-                            img_str = base64.b64encode(buffered.getvalue()).decode()
                             images_array.append(
                                 {
-                                    "data": f"data:image/png;base64,{img_str}",
+                                    "data": pil_to_data_url(pil_img),
                                     "width": pil_img.width,
                                     "height": pil_img.height,
                                 }
@@ -290,10 +286,7 @@ class LayerForgeNode:
 
                     mask_np = (input_mask.cpu().numpy() * 255).astype(np.uint8)
                     pil_mask = Image.fromarray(mask_np, "L")
-                    mask_buffered = io.BytesIO()
-                    pil_mask.save(mask_buffered, format="PNG")
-                    mask_str = base64.b64encode(mask_buffered.getvalue()).decode()
-                    input_data["input_mask"] = f"data:image/png;base64,{mask_str}"
+                    input_data["input_mask"] = pil_to_data_url(pil_mask)
                     log.debug(f"Stored input mask: {pil_mask.width}x{pil_mask.height}")
 
                 input_data["fit_on_add"] = fit_on_add
@@ -437,10 +430,7 @@ class LayerForgeNode:
 
     def get_cached_image(self):
         if self.cached_image:
-            buffered = io.BytesIO()
-            self.cached_image.save(buffered, format="PNG")
-            img_str = base64.b64encode(buffered.getvalue()).decode()
-            return f"data:image/png;base64,{img_str}"
+            return pil_to_data_url(self.cached_image)
         return None
 
 

python/routes.py

@@ -1,7 +1,5 @@
 """LayerForge HTTP and WebSocket route registration."""
 
-import base64
-import io
 import os
 import time
 
@@ -10,6 +8,7 @@ from PIL import Image
 from server import PromptServer
 
 from .matting import register_matting_routes
+from .image_serialization import file_to_data_url, pil_to_data_url
 from .node import log
 
 
@@ -99,16 +98,10 @@ def register_routes(node_class):
             }
 
             if cache_data["image"] is not None:
-                buffered = io.BytesIO()
-                cache_data["image"].save(buffered, format="PNG")
-                img_str = base64.b64encode(buffered.getvalue()).decode()
-                response_data["data"]["image"] = f"data:image/png;base64,{img_str}"
+                response_data["data"]["image"] = pil_to_data_url(cache_data["image"])
 
             if cache_data["mask"] is not None:
-                mask_buffer = io.BytesIO()
-                cache_data["mask"].save(mask_buffer, format="PNG")
-                mask_str = base64.b64encode(mask_buffer.getvalue()).decode()
-                response_data["data"]["mask"] = f"data:image/png;base64,{mask_str}"
+                response_data["data"]["mask"] = pil_to_data_url(cache_data["mask"])
 
             return web.json_response(response_data)
         except Exception as error:
@@ -122,9 +115,7 @@ def register_routes(node_class):
             latest_image_paths = node_class.get_latest_images(since_timestamp / 1000.0)
             images_data = []
             for image_path in latest_image_paths:
-                with open(image_path, "rb") as image_file:
-                    encoded_string = base64.b64encode(image_file.read()).decode("utf-8")
-                    images_data.append(f"data:image/png;base64,{encoded_string}")
+                images_data.append(file_to_data_url(image_path))
 
             return web.json_response({"success": True, "images": images_data})
         except Exception as error:
@@ -137,11 +128,9 @@ def register_routes(node_class):
         try:
             latest_image_path = node_class.get_latest_image()
             if latest_image_path:
-                with open(latest_image_path, "rb") as image_file:
-                    encoded_string = base64.b64encode(image_file.read()).decode()
                 return web.json_response({
                     "success": True,
-                    "image_data": f"data:image/png;base64,{encoded_string}",
+                    "image_data": file_to_data_url(latest_image_path),
                 })
 
             return web.json_response({
@@ -183,13 +172,10 @@ def register_routes(node_class):
                 with Image.open(file_path) as image:
                     if image.mode != "RGB":
                         image = image.convert("RGB")
-                    buffered = io.BytesIO()
-                    image.save(buffered, format="PNG")
-                    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
                     log.info(f"Successfully loaded image from path: {file_path}")
                     return web.json_response({
                         "success": True,
-                        "image_data": f"data:image/png;base64,{img_str}",
+                        "image_data": pil_to_data_url(image),
                         "width": image.width,
                         "height": image.height,
                     })