diff --git a/comfy_extras/nodes_images.py b/comfy_extras/nodes_images.py
index 7011d9c13..f864a91ed 100644
--- a/comfy_extras/nodes_images.py
+++ b/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" {escape(json.dumps(workflow))}"
+
+ packet = [
+ '',
+ '',
+ ' ',
+ f' ',
+ ]
+ if workflow_element:
+ packet.append(workflow_element)
+ packet.extend([
+ " ",
+ " ",
+ "",
+ '',
+ ])
+ 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":
diff --git a/requirements.txt b/requirements.txt
index fbdd6c14a..0efd361c2 100644
--- a/requirements.txt
+++ b/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
diff --git a/tests-unit/comfy_extras_test/nodes_images_save_test.py b/tests-unit/comfy_extras_test/nodes_images_save_test.py
new file mode 100644
index 000000000..791604002
--- /dev/null
+++ b/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": " & 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}
+ )