Add AVIF saving to Save Image Advanced

comfy_extras/nodes_images.py

@@ -14,6 +14,10 @@ import torch
 import zlib
 import comfy.utils
 from fractions import Fraction
+from io import BytesIO
+from xml.sax.saxutils import escape, quoteattr
+
+from PIL import Image as PILImage, features
 
 from server import PromptServer
 from comfy_api.latest import ComfyExtension, IO, UI
@@ -907,6 +911,7 @@ def hlg_to_linear(t: torch.Tensor) -> torch.Tensor:
 # ---------------------------------------------------------------------------
 
 _PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
+COMFYUI_XMP_NAMESPACE = "https://github.com/Comfy-Org/ComfyUI"
 
 
 def _png_chunk(chunk_type: bytes, data: bytes) -> bytes:
@@ -941,6 +946,37 @@ def inject_png_metadata(png_bytes: bytes, prompt: dict | None, extra_pnginfo: di
     return png_bytes[:ihdr_end] + b"".join(chunks) + png_bytes[ihdr_end:]
 
 
+def build_avif_xmp(prompt: dict | None, extra_pnginfo: dict | None) -> bytes | None:
+    """Build the standard XMP item used to import ComfyUI AVIF workflows."""
+    workflow = extra_pnginfo.get("workflow") if extra_pnginfo else None
+    if prompt is None and workflow is None:
+        return None
+
+    prompt_attribute = ""
+    if prompt is not None:
+        prompt_attribute = f" comfy:prompt={quoteattr(json.dumps(prompt))}"
+
+    workflow_element = ""
+    if workflow is not None:
+        workflow_element = f"      <comfy:workflow>{escape(json.dumps(workflow))}</comfy:workflow>"
+
+    packet = [
+        '<?xpacket begin="\ufeff" id="W5M0MpCehiHzreSzNTczkc9d"?>',
+        '<x:xmpmeta xmlns:x="adobe:ns:meta/">',
+        '  <rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#">',
+        f'    <rdf:Description xmlns:comfy="{COMFYUI_XMP_NAMESPACE}" rdf:about=""{prompt_attribute}>',
+    ]
+    if workflow_element:
+        packet.append(workflow_element)
+    packet.extend([
+        "    </rdf:Description>",
+        "  </rdf:RDF>",
+        "</x:xmpmeta>",
+        '<?xpacket end="w"?>',
+    ])
+    return "\n".join(packet).encode("utf-8")
+
+
 # Standard chromaticities (CIE 1931 xy) for the colorspaces this node writes.
 # Each tuple is (Rx, Ry, Gx, Gy, Bx, By, Wx, Wy). All share D65 white point.
 _CHROMATICITIES = {
@@ -1085,6 +1121,45 @@ def inject_exr_metadata(
 # Encoding
 # ---------------------------------------------------------------------------
 
+def encode_avif_image(
+    img_tensor: torch.Tensor,
+    quality: int,
+    xmp: bytes | None,
+) -> bytes:
+    """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."
+        )
+
+    if img_tensor.ndim == 2:
+        img_tensor = img_tensor.unsqueeze(-1)
+    num_channels = img_tensor.shape[-1]
+    if num_channels not in (1, 3, 4):
+        raise ValueError(
+            f"No AVIF encoder for {num_channels}-channel images: "
+            "supported channel counts are 1 (grayscale), 3 (RGB) and 4 (RGBA)."
+        )
+
+    img_np = (img_tensor * 255.0).clamp(0, 255).to(torch.uint8).cpu().numpy()
+    if num_channels == 1:
+        img_np = img_np[..., 0]
+
+    image = PILImage.fromarray(img_np)
+    output = BytesIO()
+    save_options: dict[str, int | str | bytes] = {
+        "format": "AVIF",
+        "quality": quality,
+        "subsampling": "4:2:0",
+        "range": "full",
+    }
+    if xmp is not None:
+        save_options["xmp"] = xmp
+    image.save(output, **save_options)
+    return output.getvalue()
+
 def _encode_image(
     img_tensor: torch.Tensor,
     file_format: str,
@@ -1191,6 +1266,15 @@ class SaveImageAdvanced(IO.ComfyNode):
                                 ),
                             ),
                         ]),
+                        IO.DynamicCombo.Option("avif", [
+                            IO.Int.Input(
+                                "quality",
+                                default=75,
+                                min=0,
+                                max=100,
+                                tooltip="AVIF quality. Higher values preserve more detail and produce larger files.",
+                            ),
+                        ]),
                     ],
                     tooltip="The file format in which to save the image.",
                 ),
@@ -1203,7 +1287,7 @@ class SaveImageAdvanced(IO.ComfyNode):
     @classmethod
     def execute(cls, images, filename_prefix: str, format: dict) -> IO.NodeOutput:
         file_format = format["format"]
-        bit_depth = format["bit_depth"]
+        bit_depth = format.get("bit_depth", "8-bit")
         colorspace = format.get("input_color_space", "sRGB")
 
         output_dir = folder_paths.get_output_directory()
@@ -1216,10 +1300,14 @@ class SaveImageAdvanced(IO.ComfyNode):
         prompt = cls.hidden.prompt
         extra_pnginfo = cls.hidden.extra_pnginfo
         write_metadata = not args.disable_metadata
+        avif_xmp = build_avif_xmp(prompt, extra_pnginfo) if file_format == "avif" and write_metadata else None
 
         results = []
         for batch_number, image in enumerate(images):
-            encoded = _encode_image(image, file_format, bit_depth, colorspace)
+            if file_format == "avif":
+                encoded = encode_avif_image(image, format["quality"], avif_xmp)
+            else:
+                encoded = _encode_image(image, file_format, bit_depth, colorspace)
 
             if write_metadata:
                 if file_format == "png":

requirements.txt

@@ -14,7 +14,7 @@ safetensors>=0.4.2
 aiohttp>=3.11.8
 yarl>=1.18.0
 pyyaml
-Pillow
+Pillow>=11.3.0
 scipy
 tqdm
 psutil

tests-unit/comfy_extras_test/nodes_images_save_test.py

@@ -0,0 +1,179 @@
+import json
+import sys
+from io import BytesIO
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+from xml.etree import ElementTree
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image as PILImage
+
+
+mock_nodes = MagicMock()
+mock_nodes.MAX_RESOLUTION = 16384
+mock_server = MagicMock()
+
+original_nodes = sys.modules.get("nodes")
+original_server = sys.modules.get("server")
+sys.modules["nodes"] = mock_nodes
+sys.modules["server"] = mock_server
+try:
+    from comfy_extras import nodes_images
+finally:
+    if original_nodes is None:
+        sys.modules.pop("nodes", None)
+    else:
+        sys.modules["nodes"] = original_nodes
+    if original_server is None:
+        sys.modules.pop("server", None)
+    else:
+        sys.modules["server"] = original_server
+
+
+COMFYUI_XMP_NAMESPACE = "https://github.com/Comfy-Org/ComfyUI"
+RDF_NAMESPACE = "http://www.w3.org/1999/02/22-rdf-syntax-ns#"
+
+
+def test_save_image_advanced_schema_exposes_avif_quality():
+    schema = nodes_images.SaveImageAdvanced.define_schema()
+    format_input = next(
+        input_item for input_item in schema.inputs if input_item.id == "format"
+    )
+    format_options = {option.key: option for option in format_input.options}
+
+    assert list(format_options) == ["png", "exr", "avif"]
+
+    avif_inputs = format_options["avif"].inputs
+    assert [input_item.id for input_item in avif_inputs] == ["quality"]
+    assert avif_inputs[0].default == 75
+    assert avif_inputs[0].min == 0
+    assert avif_inputs[0].max == 100
+
+    serialized_options = {
+        option["key"]: option["inputs"] for option in format_input.as_dict()["options"]
+    }
+    assert serialized_options["avif"]["required"]["quality"][:1] == ("INT",)
+
+
+def test_build_avif_xmp_escapes_and_preserves_prompt_and_workflow():
+    prompt = {"1": {"class_type": "KSampler", "inputs": {"text": "<tag> & value"}}}
+    workflow = {"nodes": [{"id": 1, "title": 'Save "AVIF"'}], "version": 1}
+
+    xmp = nodes_images.build_avif_xmp(
+        prompt, {"workflow": workflow, "ignored": {"value": 1}}
+    )
+
+    assert xmp is not None
+    root = ElementTree.fromstring(xmp)
+    description = root.find(f".//{{{RDF_NAMESPACE}}}Description")
+    workflow_element = root.find(f".//{{{COMFYUI_XMP_NAMESPACE}}}workflow")
+
+    assert description is not None
+    assert workflow_element is not None
+    assert description.attrib[f"{{{COMFYUI_XMP_NAMESPACE}}}prompt"] == json.dumps(
+        prompt
+    )
+    assert workflow_element.text == json.dumps(workflow)
+    assert b"ignored" not in xmp
+
+
+def test_encode_avif_image_passes_fixed_encoding_options(monkeypatch):
+    saved_options = {}
+    pillow_initialized = []
+
+    def fake_save(image, destination, **options):
+        saved_options.update(options)
+        destination.write(b"encoded-avif")
+
+    monkeypatch.setattr(
+        nodes_images.PILImage, "init", lambda: pillow_initialized.append(True)
+    )
+    monkeypatch.setitem(nodes_images.PILImage.SAVE, "AVIF", object())
+    monkeypatch.setattr(
+        nodes_images.features, "check", lambda feature: feature == "avif"
+    )
+    monkeypatch.setattr(PILImage.Image, "save", fake_save)
+
+    image = torch.zeros((8, 12, 3), dtype=torch.float32)
+    encoded = nodes_images.encode_avif_image(image, quality=63, xmp=b"xmp-packet")
+
+    assert encoded == b"encoded-avif"
+    assert pillow_initialized == [True]
+    assert saved_options == {
+        "format": "AVIF",
+        "quality": 63,
+        "subsampling": "4:2:0",
+        "range": "full",
+        "xmp": b"xmp-packet",
+    }
+
+
+def test_encode_avif_image_round_trips_rgba_and_xmp():
+    prompt = {"1": {"class_type": "KSampler"}}
+    workflow = {"nodes": [], "version": 1}
+    xmp = nodes_images.build_avif_xmp(prompt, {"workflow": workflow})
+    image = torch.rand((12, 16, 4), dtype=torch.float32)
+
+    encoded = nodes_images.encode_avif_image(image, quality=85, xmp=xmp)
+
+    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 == "RGBA"
+        assert saved_image.info["xmp"] == xmp
+
+
+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"):
+        nodes_images.encode_avif_image(torch.zeros((8, 8, 3)), quality=75, xmp=None)
+
+
+@pytest.mark.parametrize("disable_metadata", [False, True])
+def test_save_image_advanced_writes_avif_batch_and_respects_metadata_setting(
+    monkeypatch, tmp_path, disable_metadata
+):
+    prompt = {"1": {"class_type": "KSampler"}}
+    workflow = {"nodes": [], "version": 1}
+    images = torch.from_numpy(
+        np.random.default_rng(7).random((2, 12, 16, 3), dtype=np.float32)
+    )
+
+    monkeypatch.setattr(nodes_images.args, "disable_metadata", disable_metadata)
+    monkeypatch.setattr(
+        nodes_images.folder_paths, "get_output_directory", lambda: str(tmp_path)
+    )
+    monkeypatch.setattr(
+        nodes_images.folder_paths,
+        "get_save_image_path",
+        lambda *args: (str(tmp_path), "ComfyUI_%batch_num%", 7, "", "ComfyUI"),
+    )
+    monkeypatch.setattr(
+        nodes_images.SaveImageAdvanced,
+        "hidden",
+        SimpleNamespace(prompt=prompt, extra_pnginfo={"workflow": workflow}),
+    )
+
+    output = nodes_images.SaveImageAdvanced.execute(
+        images,
+        filename_prefix="ComfyUI",
+        format={"format": "avif", "quality": 82},
+    )
+
+    expected_filenames = ["ComfyUI_0_00007.avif", "ComfyUI_1_00008.avif"]
+    assert output[0] is images
+    assert [result["filename"] for result in output.ui["images"]] == expected_filenames
+
+    for filename in expected_filenames:
+        with PILImage.open(tmp_path / filename) as saved_image:
+            saved_image.load()
+            if disable_metadata:
+                assert "xmp" not in saved_image.info
+            else:
+                assert saved_image.info["xmp"] == nodes_images.build_avif_xmp(
+                    prompt, {"workflow": workflow}
+                )