From ce8124b1bd0209e9d416786830ded06a2800a7bf Mon Sep 17 00:00:00 2001 From: Dixing Xu Date: Fri, 10 Jul 2026 16:47:36 +0100 Subject: [PATCH] perf(image-files): select RLE for high-entropy PNGs Fixes #6597 --- .../services/image_files/image_files_disk.py | 48 +++++++++- .../image_files/test_image_files_disk.py | 96 ++++++++++++++++++- 2 files changed, 141 insertions(+), 3 deletions(-) diff --git a/invokeai/app/services/image_files/image_files_disk.py b/invokeai/app/services/image_files/image_files_disk.py index c54f767b164..929775a07f0 100644 --- a/invokeai/app/services/image_files/image_files_disk.py +++ b/invokeai/app/services/image_files/image_files_disk.py @@ -1,4 +1,6 @@ # Copyright (c) 2022 Kyle Schouviller (https://github.com/kyle0654) and the InvokeAI Team +import io +import zlib from pathlib import Path from queue import Queue from typing import Optional, Union @@ -15,6 +17,46 @@ from invokeai.app.services.invoker import Invoker from invokeai.app.util.thumbnails import get_thumbnail_name, make_thumbnail +_PNG_RLE_MIN_PIXELS = 512 * 512 +_PNG_RLE_SAMPLE_TILE = 32 +_PNG_RLE_MIN_RAW_SIZE_PERCENT = 30 +_PNG_RLE_MAX_SAMPLE_SIZE_PERCENT = 102 + + +def _get_png_size(image: PILImageType, compress_type: Optional[int] = None) -> int: + output = io.BytesIO() + options = {"compress_level": 1} + if compress_type is not None: + options["compress_type"] = compress_type + image.save(output, "PNG", **options) + return output.tell() + + +def _should_use_png_rle(image: PILImageType) -> bool: + if image.mode not in {"RGB", "RGBA"} or image.width * image.height < _PNG_RLE_MIN_PIXELS: + return False + + # Native-resolution tiles distinguish high-entropy data from filter-friendly structured images. + tile_width = min(_PNG_RLE_SAMPLE_TILE, image.width) + tile_height = min(_PNG_RLE_SAMPLE_TILE, image.height) + x_positions = (0, (image.width - tile_width) // 2, image.width - tile_width) + y_positions = (0, (image.height - tile_height) // 2, image.height - tile_height) + sample = Image.new(image.mode, (tile_width * 3, tile_height * 3)) + for row, y in enumerate(y_positions): + for column, x in enumerate(x_positions): + with image.crop((x, y, x + tile_width, y + tile_height)) as tile: + sample.paste(tile, (column * tile_width, row * tile_height)) + + try: + raw = sample.tobytes() + if len(zlib.compress(raw, level=1)) * 100 < len(raw) * _PNG_RLE_MIN_RAW_SIZE_PERCENT: + return False + default_size = _get_png_size(sample) + rle_size = _get_png_size(sample, zlib.Z_RLE) + return rle_size * 100 <= default_size * _PNG_RLE_MAX_SAMPLE_SIZE_PERCENT + finally: + sample.close() + class DiskImageFileStorage(ImageFileStorageBase): """Stores images on disk""" @@ -90,11 +132,15 @@ def save( # When saving the image, the image object's info field is not populated. We need to set it image.info = info_dict + compress_level = self.__invoker.services.configuration.pil_compress_level + save_options = {"compress_level": compress_level} + if compress_level == 1 and _should_use_png_rle(image): + save_options["compress_type"] = zlib.Z_RLE image.save( image_path, "PNG", pnginfo=pnginfo, - compress_level=self.__invoker.services.configuration.pil_compress_level, + **save_options, ) thumbnail_path = self.get_path(image_name, thumbnail=True, image_subfolder=image_subfolder) diff --git a/tests/app/services/image_files/test_image_files_disk.py b/tests/app/services/image_files/test_image_files_disk.py index 58c1a78635f..9a819568bc0 100644 --- a/tests/app/services/image_files/test_image_files_disk.py +++ b/tests/app/services/image_files/test_image_files_disk.py @@ -1,11 +1,13 @@ +import hashlib import platform +import zlib from pathlib import Path -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest from PIL import Image -from invokeai.app.services.image_files.image_files_disk import DiskImageFileStorage +from invokeai.app.services.image_files.image_files_disk import DiskImageFileStorage, _should_use_png_rle from invokeai.app.util.thumbnails import get_thumbnail_name @@ -64,6 +66,96 @@ def test_image_paths_relative_to_storage_dir(tmp_path: Path): assert path.is_relative_to(tmp_path) +@pytest.mark.parametrize( + ("compress_level", "expected_compress_type"), + [(0, None), (1, zlib.Z_RLE), (7, None)], +) +def test_save_uses_rle_only_for_compression_level_one( + tmp_path: Path, compress_level: int, expected_compress_type: int | None +): + storage = DiskImageFileStorage(tmp_path) + mock_invoker = MagicMock() + mock_invoker.services.configuration.pil_compress_level = compress_level + storage._DiskImageFileStorage__invoker = mock_invoker # type: ignore + + with ( + patch("invokeai.app.services.image_files.image_files_disk._should_use_png_rle", return_value=True), + patch.object(Image.Image, "save", autospec=True) as save_mock, + ): + storage.save(image=Image.new("RGBA", (32, 32)), image_name="test.png") + + png_calls = [call for call in save_mock.call_args_list if len(call.args) > 2 and call.args[2] == "PNG"] + assert len(png_calls) == 1 + assert png_calls[0].kwargs["compress_level"] == compress_level + if expected_compress_type is None: + assert "compress_type" not in png_calls[0].kwargs + else: + assert png_calls[0].kwargs["compress_type"] == expected_compress_type + + +def test_png_rle_probe_rejects_structured_images(): + entropy = Image.frombytes("RGB", (512, 512), hashlib.shake_256(b"png-rle-test").digest(512 * 512 * 3)) + gradient = Image.linear_gradient("L").resize((512, 512)).convert("RGB") + + assert _should_use_png_rle(entropy) + assert not _should_use_png_rle(gradient) + + entropy.close() + gradient.close() + + +def _make_round_trip_image(mode: str) -> Image.Image: + image = Image.new(mode, (4, 4)) + if mode == "P": + palette = [component for index in range(256) for component in (index, 255 - index, index // 2, index)] + image.putpalette(palette, rawmode="RGBA") + image.putdata(range(16)) + else: + values = { + "1": [0, 1], + "L": [0, 255], + "LA": [(17, 0), (201, 255)], + "RGB": [(1, 2, 3), (251, 252, 253)], + "RGBA": [(1, 2, 3, 0), (251, 252, 253, 255)], + "I;16": [0, 65535], + } + image.putdata(values[mode] * 8) + return image + + +@pytest.mark.parametrize("mode", ["1", "L", "LA", "P", "RGB", "RGBA", "I;16"]) +def test_level_one_png_round_trip_from_disk(tmp_path: Path, mode: str): + storage = DiskImageFileStorage(tmp_path) + mock_invoker = MagicMock() + mock_invoker.services.configuration.pil_compress_level = 1 + storage._DiskImageFileStorage__invoker = mock_invoker # type: ignore + + image = _make_round_trip_image(mode) + expected_bytes = image.tobytes() + expected_rgba = image.convert("RGBA").tobytes() if mode == "P" else None + metadata = f'{{"mode":"{mode}"}}' + image_name = f"round-trip-{mode.replace(';', '-')}.png" + + with patch("invokeai.app.services.image_files.image_files_disk._should_use_png_rle", return_value=True): + storage.save(image=image, image_name=image_name, metadata=metadata) + image_path = storage.get_path(image_name) + storage.evict_cache_paths([image_path]) + + with Image.open(image_path) as loaded: + loaded.load() + assert loaded.format == "PNG" + assert loaded.mode == mode + assert loaded.tobytes() == expected_bytes + assert loaded.info["invokeai_metadata"] == metadata + if mode in {"LA", "RGBA"}: + assert loaded.getchannel("A").tobytes() == image.getchannel("A").tobytes() + if mode == "P": + assert loaded.info["transparency"] == bytes(range(256)) + assert loaded.convert("RGBA").tobytes() == expected_rgba + + image.close() + + # ── Subfolder validation tests (Point 1) ──