diff --git a/comfy_extras/nodes_images.py b/comfy_extras/nodes_images.py index 7011d9c13..a47de69f9 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,43 @@ 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 a Pillow build with AVIF support.") + + 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 +1264,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 +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["bit_depth"] colorspace = format.get("input_color_space", "sRGB") output_dir = folder_paths.get_output_directory() @@ -1216,10 +1297,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, 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..eadc4aeb8 --- /dev/null +++ b/tests-unit/comfy_extras_test/nodes_images_save_test.py @@ -0,0 +1,210 @@ +import importlib +import json +import sys +from io import BytesIO +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", 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: + nodes_images = importlib.import_module("comfy_extras.nodes_images") +finally: + if original_nodes is MISSING: + sys.modules.pop("nodes", None) + else: + sys.modules["nodes"] = original_nodes + 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" +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 + + +@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 build 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} + )