From f2aa3b74ae7c7250700e8e650722da88e90cf06c Mon Sep 17 00:00:00 2001 From: Austin Traver Date: Fri, 17 Jul 2026 10:55:20 -0700 Subject: [PATCH] Address AVIF review feedback --- comfy_extras/nodes_images.py | 9 ++-- .../nodes_images_save_test.py | 43 ++++++++++++++++--- 2 files changed, 40 insertions(+), 12 deletions(-) diff --git a/comfy_extras/nodes_images.py b/comfy_extras/nodes_images.py index f864a91ed..a47de69f9 100644 --- a/comfy_extras/nodes_images.py +++ b/comfy_extras/nodes_images.py @@ -1129,10 +1129,7 @@ def encode_avif_image( """Encode one sRGB image tensor as an 8-bit AVIF with optional XMP.""" PILImage.init() if not features.check("avif") or "AVIF" not in PILImage.SAVE: - raise RuntimeError( - "Saving AVIF images requires Pillow built with AVIF support. " - "Official Pillow 11.3.0 and newer wheels include it." - ) + raise RuntimeError("Saving AVIF images requires a Pillow build with AVIF support.") if img_tensor.ndim == 2: img_tensor = img_tensor.unsqueeze(-1) @@ -1160,6 +1157,7 @@ def encode_avif_image( image.save(output, **save_options) return output.getvalue() + def _encode_image( img_tensor: torch.Tensor, file_format: str, @@ -1287,7 +1285,6 @@ class SaveImageAdvanced(IO.ComfyNode): @classmethod def execute(cls, images, filename_prefix: str, format: dict) -> IO.NodeOutput: file_format = format["format"] - bit_depth = format.get("bit_depth", "8-bit") colorspace = format.get("input_color_space", "sRGB") output_dir = folder_paths.get_output_directory() @@ -1307,7 +1304,7 @@ class SaveImageAdvanced(IO.ComfyNode): if file_format == "avif": encoded = encode_avif_image(image, format["quality"], avif_xmp) else: - encoded = _encode_image(image, file_format, bit_depth, colorspace) + encoded = _encode_image(image, file_format, format["bit_depth"], colorspace) if write_metadata: if file_format == "png": diff --git a/tests-unit/comfy_extras_test/nodes_images_save_test.py b/tests-unit/comfy_extras_test/nodes_images_save_test.py index 791604002..eadc4aeb8 100644 --- a/tests-unit/comfy_extras_test/nodes_images_save_test.py +++ b/tests-unit/comfy_extras_test/nodes_images_save_test.py @@ -1,3 +1,4 @@ +import importlib import json import sys from io import BytesIO @@ -5,31 +6,48 @@ from types import SimpleNamespace from unittest.mock import MagicMock from xml.etree import ElementTree +import comfy_extras import numpy as np import pytest import torch from PIL import Image as PILImage +MISSING = object() + mock_nodes = MagicMock() mock_nodes.MAX_RESOLUTION = 16384 mock_server = MagicMock() -original_nodes = sys.modules.get("nodes") -original_server = sys.modules.get("server") +original_nodes = sys.modules.get("nodes", MISSING) +original_server = sys.modules.get("server", MISSING) +original_nodes_images = sys.modules.get("comfy_extras.nodes_images", MISSING) +original_nodes_images_attr = getattr(comfy_extras, "nodes_images", MISSING) sys.modules["nodes"] = mock_nodes sys.modules["server"] = mock_server +sys.modules.pop("comfy_extras.nodes_images", None) +if original_nodes_images_attr is not MISSING: + delattr(comfy_extras, "nodes_images") try: - from comfy_extras import nodes_images + nodes_images = importlib.import_module("comfy_extras.nodes_images") finally: - if original_nodes is None: + if original_nodes is MISSING: sys.modules.pop("nodes", None) else: sys.modules["nodes"] = original_nodes - if original_server is None: + if original_server is MISSING: sys.modules.pop("server", None) else: sys.modules["server"] = original_server + if original_nodes_images is MISSING: + sys.modules.pop("comfy_extras.nodes_images", None) + else: + sys.modules["comfy_extras.nodes_images"] = original_nodes_images + if original_nodes_images_attr is MISSING: + if hasattr(comfy_extras, "nodes_images"): + delattr(comfy_extras, "nodes_images") + else: + setattr(comfy_extras, "nodes_images", original_nodes_images_attr) COMFYUI_XMP_NAMESPACE = "https://github.com/Comfy-Org/ComfyUI" @@ -126,10 +144,23 @@ def test_encode_avif_image_round_trips_rgba_and_xmp(): assert saved_image.info["xmp"] == xmp +@pytest.mark.parametrize("image_shape", [(12, 16), (12, 16, 1)]) +def test_encode_avif_image_round_trips_grayscale(image_shape): + image = torch.rand(image_shape, dtype=torch.float32) + + encoded = nodes_images.encode_avif_image(image, quality=85, xmp=None) + + with PILImage.open(BytesIO(encoded)) as saved_image: + saved_image.load() + assert saved_image.format == "AVIF" + assert saved_image.size == (16, 12) + assert saved_image.mode == "RGB" + + def test_encode_avif_image_reports_missing_pillow_support(monkeypatch): monkeypatch.setattr(nodes_images.features, "check", lambda feature: False) - with pytest.raises(RuntimeError, match="Pillow built with AVIF support"): + with pytest.raises(RuntimeError, match="Pillow build with AVIF support"): nodes_images.encode_avif_image(torch.zeros((8, 8, 3)), quality=75, xmp=None)