Address AVIF review feedback

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":

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)