mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-08-28 11:57:41 +08:00
Support avif in Save Image Advanced node. (#15891)
This commit is contained in:
comfyanonymous
@@ -11,10 +11,13 @@ import numpy as np
|
||||
import struct
|
||||
import torch
|
||||
import logging
|
||||
import tempfile
|
||||
|
||||
import zlib
|
||||
import comfy.utils
|
||||
from av.video.reformatter import ColorPrimaries, ColorRange, ColorTrc
|
||||
from fractions import Fraction
|
||||
from PIL.Image import Exif
|
||||
|
||||
from server import PromptServer
|
||||
from comfy_api.latest import ComfyExtension, IO, UI
|
||||
@@ -1002,6 +1005,12 @@ _FORMAT_SPECS = {
|
||||
("exr", "32-bit float", 4): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "gbrapf32le", "stream_fmt": "gbrapf32le"},
|
||||
}
|
||||
|
||||
_AVIF_COLOR_PROPERTIES = {
|
||||
"sRGB": (ColorPrimaries.BT709, ColorTrc.IEC61966_2_1, 1),
|
||||
"HDR": (ColorPrimaries.BT2020, ColorTrc.ARIB_STD_B67, 9),
|
||||
"HDR PQ": (ColorPrimaries.BT2020, ColorTrc.SMPTE2084, 9),
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Color transforms
|
||||
@@ -1225,6 +1234,236 @@ def inject_exr_metadata(
|
||||
)
|
||||
|
||||
|
||||
def _bmff_box(box_type: bytes, payload: bytes) -> bytes:
|
||||
size = 8 + len(payload)
|
||||
if size > 0xFFFFFFFF:
|
||||
raise ValueError("AVIF metadata box is too large.")
|
||||
return struct.pack(">I4s", size, box_type) + payload
|
||||
|
||||
|
||||
def _bmff_boxes(data: bytes, start: int, end: int) -> list[tuple[int, int, bytes, int]]:
|
||||
boxes = []
|
||||
pos = start
|
||||
while pos < end:
|
||||
if pos + 8 > end:
|
||||
raise ValueError("Invalid AVIF box structure.")
|
||||
size, box_type = struct.unpack_from(">I4s", data, pos)
|
||||
header_size = 8
|
||||
if size == 1:
|
||||
if pos + 16 > end:
|
||||
raise ValueError("Invalid AVIF extended-size box.")
|
||||
size = struct.unpack_from(">Q", data, pos + 8)[0]
|
||||
header_size = 16
|
||||
elif size == 0:
|
||||
size = end - pos
|
||||
if size < header_size or pos + size > end:
|
||||
raise ValueError("Invalid AVIF box size.")
|
||||
boxes.append((pos, size, box_type, header_size))
|
||||
pos += size
|
||||
return boxes
|
||||
|
||||
|
||||
def _avif_exif(metadata: dict) -> bytes:
|
||||
exif = Exif()
|
||||
if "prompt" in metadata:
|
||||
exif[0x0110] = f"prompt:{json.dumps(metadata['prompt'])}"
|
||||
next_tag = 0x010F
|
||||
for key, value in metadata.items():
|
||||
if key == "prompt":
|
||||
continue
|
||||
exif[next_tag] = f"{key}:{json.dumps(value)}"
|
||||
next_tag -= 1
|
||||
return b"\x00\x00\x00\x00" + exif.tobytes()[6:]
|
||||
|
||||
|
||||
def _add_avif_exif_item(meta: bytes, exif_offset: int, exif_length: int, offset_delta: int) -> bytes:
|
||||
if meta[4:8] != b"meta" or len(meta) < 12:
|
||||
raise ValueError("AVIF metadata requires a valid meta box.")
|
||||
children = _bmff_boxes(meta, 12, len(meta))
|
||||
child_by_type = {box_type: (pos, size, header_size) for pos, size, box_type, header_size in children}
|
||||
if not all(box_type in child_by_type for box_type in (b"pitm", b"iloc", b"iinf")):
|
||||
raise ValueError("AVIF metadata boxes are incomplete.")
|
||||
|
||||
pitm_pos, pitm_size, _ = child_by_type[b"pitm"]
|
||||
pitm = meta[pitm_pos:pitm_pos + pitm_size]
|
||||
if pitm[8] != 0:
|
||||
raise ValueError("Unsupported AVIF primary-item format.")
|
||||
primary_item_id = struct.unpack_from(">H", pitm, 12)[0]
|
||||
|
||||
iloc_pos, iloc_size, _ = child_by_type[b"iloc"]
|
||||
iloc = bytearray(meta[iloc_pos:iloc_pos + iloc_size])
|
||||
if iloc[8] != 0 or iloc[12] != 0x44 or iloc[13] != 0:
|
||||
raise ValueError("Unsupported AVIF item-location format.")
|
||||
item_count = struct.unpack_from(">H", iloc, 14)[0]
|
||||
cursor = 16
|
||||
item_ids = []
|
||||
for _ in range(item_count):
|
||||
item_id, _, extent_count = struct.unpack_from(">HHH", iloc, cursor)
|
||||
cursor += 6
|
||||
item_ids.append(item_id)
|
||||
for _ in range(extent_count):
|
||||
extent_offset = struct.unpack_from(">I", iloc, cursor)[0]
|
||||
struct.pack_into(">I", iloc, cursor, extent_offset + offset_delta)
|
||||
cursor += 8
|
||||
if cursor != len(iloc):
|
||||
raise ValueError("Unsupported AVIF item-location entries.")
|
||||
exif_item_id = max(item_ids) + 1
|
||||
if exif_item_id > 0xFFFF or exif_offset > 0xFFFFFFFF or exif_length > 0xFFFFFFFF:
|
||||
raise ValueError("AVIF metadata exceeds 32-bit item limits.")
|
||||
struct.pack_into(">H", iloc, 14, item_count + 1)
|
||||
iloc.extend(struct.pack(">HHHII", exif_item_id, 0, 1, exif_offset, exif_length))
|
||||
struct.pack_into(">I", iloc, 0, len(iloc))
|
||||
|
||||
iinf_pos, iinf_size, _ = child_by_type[b"iinf"]
|
||||
iinf = bytearray(meta[iinf_pos:iinf_pos + iinf_size])
|
||||
if iinf[8] != 0:
|
||||
raise ValueError("Unsupported AVIF item-information format.")
|
||||
iinf_count = struct.unpack_from(">H", iinf, 12)[0]
|
||||
struct.pack_into(">H", iinf, 12, iinf_count + 1)
|
||||
infe_payload = b"\x02\x00\x00\x00" + struct.pack(">HH4s", exif_item_id, 0, b"Exif") + b"\x00"
|
||||
iinf.extend(_bmff_box(b"infe", infe_payload))
|
||||
struct.pack_into(">I", iinf, 0, len(iinf))
|
||||
|
||||
cdsc = _bmff_box(b"cdsc", struct.pack(">HHH", exif_item_id, 1, primary_item_id))
|
||||
if b"iref" in child_by_type:
|
||||
iref_pos, iref_size, _ = child_by_type[b"iref"]
|
||||
iref = bytearray(meta[iref_pos:iref_pos + iref_size])
|
||||
if iref[8] != 0:
|
||||
raise ValueError("Unsupported AVIF item-reference format.")
|
||||
iref.extend(cdsc)
|
||||
struct.pack_into(">I", iref, 0, len(iref))
|
||||
else:
|
||||
iref = _bmff_box(b"iref", b"\x00\x00\x00\x00" + cdsc)
|
||||
|
||||
output = bytearray(meta[:12])
|
||||
for pos, size, box_type, _ in children:
|
||||
if box_type == b"iloc":
|
||||
output.extend(iloc)
|
||||
elif box_type == b"iinf":
|
||||
output.extend(iinf)
|
||||
if b"iref" not in child_by_type:
|
||||
output.extend(iref)
|
||||
elif box_type == b"iref":
|
||||
output.extend(iref)
|
||||
else:
|
||||
output.extend(meta[pos:pos + size])
|
||||
struct.pack_into(">I", output, 0, len(output))
|
||||
return bytes(output)
|
||||
|
||||
|
||||
def _adjust_avif_chunk_offsets(moov: bytes, offset_delta: int) -> bytes:
|
||||
output = bytearray(moov)
|
||||
containers = {b"moov", b"trak", b"mdia", b"minf", b"stbl"}
|
||||
|
||||
def adjust(start: int, end: int) -> None:
|
||||
for pos, size, box_type, header_size in _bmff_boxes(output, start, end):
|
||||
if box_type in containers:
|
||||
adjust(pos + header_size, pos + size)
|
||||
elif box_type in (b"stco", b"co64"):
|
||||
entry_size = 4 if box_type == b"stco" else 8
|
||||
entry_count = struct.unpack_from(">I", output, pos + header_size + 4)[0]
|
||||
cursor = pos + header_size + 8
|
||||
if cursor + entry_count * entry_size != pos + size:
|
||||
raise ValueError("Invalid AVIF chunk-offset table.")
|
||||
value_format = ">I" if entry_size == 4 else ">Q"
|
||||
for _ in range(entry_count):
|
||||
value = struct.unpack_from(value_format, output, cursor)[0]
|
||||
struct.pack_into(value_format, output, cursor, value + offset_delta)
|
||||
cursor += entry_size
|
||||
|
||||
adjust(0, len(output))
|
||||
return bytes(output)
|
||||
|
||||
|
||||
def _copy_file_bytes(source, destination, size: int) -> None:
|
||||
remaining = size
|
||||
while remaining:
|
||||
chunk = source.read(min(remaining, 1024 * 1024))
|
||||
if not chunk:
|
||||
raise ValueError("Unexpected end of AVIF file.")
|
||||
destination.write(chunk)
|
||||
remaining -= len(chunk)
|
||||
|
||||
|
||||
def inject_avif_metadata(path: str, metadata: dict) -> None:
|
||||
"""Add a ComfyUI-compatible EXIF item without loading media data into memory."""
|
||||
if not metadata:
|
||||
return
|
||||
exif = _avif_exif(metadata)
|
||||
file_size = os.path.getsize(path)
|
||||
top_level_boxes = []
|
||||
with open(path, "rb") as source:
|
||||
pos = 0
|
||||
while pos < file_size:
|
||||
source.seek(pos)
|
||||
header = source.read(8)
|
||||
if len(header) != 8:
|
||||
raise ValueError("Invalid AVIF box header.")
|
||||
size, box_type = struct.unpack(">I4s", header)
|
||||
header_size = 8
|
||||
size_field = size
|
||||
if size == 1:
|
||||
extended_size = source.read(8)
|
||||
if len(extended_size) != 8:
|
||||
raise ValueError("Invalid AVIF extended-size box.")
|
||||
size = struct.unpack(">Q", extended_size)[0]
|
||||
header_size = 16
|
||||
elif size == 0:
|
||||
size = file_size - pos
|
||||
if size < header_size or pos + size > file_size:
|
||||
raise ValueError("Invalid AVIF top-level box size.")
|
||||
top_level_boxes.append((pos, size, box_type, header_size, size_field))
|
||||
pos += size
|
||||
|
||||
meta_box = next((box for box in top_level_boxes if box[2] == b"meta"), None)
|
||||
mdat_box = next((box for box in top_level_boxes if box[2] == b"mdat"), None)
|
||||
if meta_box is None or mdat_box is None or mdat_box[0] + mdat_box[1] != file_size or meta_box[0] > mdat_box[0]:
|
||||
raise ValueError("Unsupported AVIF file layout for metadata.")
|
||||
source.seek(meta_box[0])
|
||||
meta = source.read(meta_box[1])
|
||||
|
||||
provisional_meta = _add_avif_exif_item(meta, 0, len(exif), 0)
|
||||
offset_delta = len(provisional_meta) - len(meta)
|
||||
exif_offset = mdat_box[0] + mdat_box[1] + offset_delta
|
||||
updated_meta = _add_avif_exif_item(meta, exif_offset, len(exif), offset_delta)
|
||||
|
||||
fd, temp_path = tempfile.mkstemp(prefix=f".{os.path.basename(path)}.", dir=os.path.dirname(path) or ".")
|
||||
try:
|
||||
with os.fdopen(fd, "wb") as destination, open(path, "rb") as source:
|
||||
for pos, size, box_type, header_size, size_field in top_level_boxes:
|
||||
source.seek(pos)
|
||||
if box_type == b"ftyp":
|
||||
ftyp = bytearray(source.read(size))
|
||||
# The frontend metadata reader recognizes the avif major brand; avis remains a compatible sequence brand.
|
||||
if ftyp[8:12] == b"avis" and b"avif" in ftyp[16:]:
|
||||
ftyp[8:12] = b"avif"
|
||||
destination.write(ftyp)
|
||||
elif box_type == b"meta":
|
||||
destination.write(updated_meta)
|
||||
elif box_type == b"moov":
|
||||
destination.write(_adjust_avif_chunk_offsets(source.read(size), offset_delta))
|
||||
elif box_type == b"mdat":
|
||||
new_size = size + len(exif)
|
||||
if size_field == 1:
|
||||
destination.write(struct.pack(">I4sQ", 1, b"mdat", new_size))
|
||||
elif size_field == 0:
|
||||
destination.write(struct.pack(">I4s", 0, b"mdat"))
|
||||
elif new_size <= 0xFFFFFFFF:
|
||||
destination.write(struct.pack(">I4s", new_size, b"mdat"))
|
||||
else:
|
||||
raise ValueError("AVIF media-data box exceeds its 32-bit size field.")
|
||||
source.seek(pos + header_size)
|
||||
_copy_file_bytes(source, destination, size - header_size)
|
||||
destination.write(exif)
|
||||
else:
|
||||
_copy_file_bytes(source, destination, size)
|
||||
os.chmod(temp_path, os.stat(path).st_mode)
|
||||
os.replace(temp_path, path)
|
||||
finally:
|
||||
if os.path.exists(temp_path):
|
||||
os.unlink(temp_path)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Encoding
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -1292,6 +1531,69 @@ def _encode_image(
|
||||
return b"".join(bytes(p) for p in packets)
|
||||
|
||||
|
||||
def _set_avif_color_properties(target, colorspace: str) -> None:
|
||||
color_primaries, color_trc, yuv_colorspace = _AVIF_COLOR_PROPERTIES[colorspace]
|
||||
target.color_primaries = color_primaries
|
||||
target.color_trc = color_trc
|
||||
target.colorspace = yuv_colorspace
|
||||
target.color_range = ColorRange.MPEG
|
||||
|
||||
|
||||
def _avif_frame(image: torch.Tensor, bit_depth: str, colorspace: str, pixel_format: str) -> av.VideoFrame:
|
||||
if image.ndim == 2:
|
||||
image = image.unsqueeze(-1)
|
||||
num_channels = image.shape[-1]
|
||||
if num_channels not in (1, 3):
|
||||
raise ValueError("AVIF saving supports 1-channel grayscale and 3-channel RGB images; PyAV's SVT-AV1 encoder does not support alpha.")
|
||||
|
||||
if bit_depth == "10-bit YUV420":
|
||||
image_np = (image * 65535.0).clamp(0, 65535).to(torch.int32).cpu().numpy().astype(np.uint16)
|
||||
frame_format = "gray16le" if num_channels == 1 else "rgb48le"
|
||||
else:
|
||||
image_np = (image * 255.0).clamp(0, 255).to(torch.uint8).cpu().numpy()
|
||||
frame_format = "gray" if num_channels == 1 else "rgb24"
|
||||
if num_channels == 1:
|
||||
image_np = image_np[..., 0]
|
||||
|
||||
frame = av.VideoFrame.from_ndarray(image_np, format=frame_format)
|
||||
frame = frame.reformat(format=pixel_format, dst_colorspace=_AVIF_COLOR_PROPERTIES[colorspace][2])
|
||||
_set_avif_color_properties(frame, colorspace)
|
||||
return frame
|
||||
|
||||
|
||||
def _save_avif(
|
||||
images: torch.Tensor,
|
||||
output_path: str,
|
||||
bit_depth: str,
|
||||
colorspace: str,
|
||||
crf: int,
|
||||
fps: float = 1.0,
|
||||
loop_count: int | None = None,
|
||||
metadata: dict | None = None,
|
||||
) -> None:
|
||||
if bit_depth == "auto":
|
||||
bit_depth = "10-bit YUV420" if colorspace in ("HDR", "HDR PQ") else "8-bit YUV420"
|
||||
pixel_format = "yuv420p10le" if bit_depth == "10-bit YUV420" else "yuv420p"
|
||||
options = {"loop": str(loop_count)} if loop_count is not None else None
|
||||
|
||||
with av.open(output_path, mode="w", format="avif", options=options) as container:
|
||||
frame_rate = Fraction(round(fps * 1000), 1000)
|
||||
stream = container.add_stream("libsvtav1", rate=frame_rate)
|
||||
stream.width = images.shape[2]
|
||||
stream.height = images.shape[1]
|
||||
stream.pix_fmt = pixel_format
|
||||
stream.options = {"crf": str(crf), "preset": "8"}
|
||||
_set_avif_color_properties(stream.codec_context, colorspace)
|
||||
|
||||
for image in images:
|
||||
for packet in stream.encode(_avif_frame(image, bit_depth, colorspace, pixel_format)):
|
||||
container.mux(packet)
|
||||
for packet in stream.encode(None):
|
||||
container.mux(packet)
|
||||
if metadata:
|
||||
inject_avif_metadata(output_path, metadata)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Node
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -1335,6 +1637,48 @@ class SaveImageAdvanced(IO.ComfyNode):
|
||||
),
|
||||
),
|
||||
]),
|
||||
IO.DynamicCombo.Option("avif", [
|
||||
IO.Combo.Input(
|
||||
"bit_depth",
|
||||
options=["auto", "8-bit YUV420", "10-bit YUV420"],
|
||||
default="auto",
|
||||
advanced=True,
|
||||
tooltip="Auto uses 8-bit YUV420 for sRGB and 10-bit YUV420 for HDR.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"input_color_space",
|
||||
options=["sRGB", "HDR", "HDR PQ"],
|
||||
default="sRGB",
|
||||
advanced=True,
|
||||
tooltip="Colorspace of the input images. HDR selects BT.2020/HLG and HDR PQ selects BT.2020/PQ.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"crf",
|
||||
default=18,
|
||||
min=1,
|
||||
max=63,
|
||||
advanced=True,
|
||||
tooltip="Lower values produce higher quality and larger files.",
|
||||
),
|
||||
IO.DynamicCombo.Input(
|
||||
"save_mode",
|
||||
display_name="save mode",
|
||||
options=[
|
||||
IO.DynamicCombo.Option("still images", []),
|
||||
IO.DynamicCombo.Option("animated", [
|
||||
IO.Float.Input("fps", default=6.0, min=0.01, max=1000.0, step=0.01),
|
||||
IO.Int.Input(
|
||||
"loop_count",
|
||||
default=0,
|
||||
min=0,
|
||||
max=1000,
|
||||
advanced=True,
|
||||
tooltip="Number of times to loop the animation. 0 loops forever.",
|
||||
),
|
||||
]),
|
||||
],
|
||||
),
|
||||
]),
|
||||
],
|
||||
tooltip="The file format in which to save the image.",
|
||||
),
|
||||
@@ -1362,6 +1706,37 @@ class SaveImageAdvanced(IO.ComfyNode):
|
||||
write_metadata = not args.disable_metadata
|
||||
|
||||
results = []
|
||||
if file_format == "avif":
|
||||
metadata = None
|
||||
if write_metadata:
|
||||
metadata = {}
|
||||
if prompt is not None:
|
||||
metadata["prompt"] = prompt
|
||||
if extra_pnginfo:
|
||||
metadata.update(extra_pnginfo)
|
||||
save_mode = format["save_mode"]
|
||||
animated = save_mode["save_mode"] == "animated"
|
||||
batches = [images] if animated else images.unsqueeze(1)
|
||||
for batch_number, batch in enumerate(batches):
|
||||
name = filename.replace("%batch_num%", str(batch_number))
|
||||
file = f"{name}_{counter:05}.avif"
|
||||
_save_avif(
|
||||
batch,
|
||||
os.path.join(full_output_folder, file),
|
||||
bit_depth,
|
||||
colorspace,
|
||||
format["crf"],
|
||||
fps=save_mode.get("fps", 1.0),
|
||||
loop_count=save_mode.get("loop_count") if animated else None,
|
||||
metadata=metadata,
|
||||
)
|
||||
results.append({"filename": file, "subfolder": subfolder, "type": "output"})
|
||||
counter += 1
|
||||
ui = {"images": results}
|
||||
if animated and len(images) > 1:
|
||||
ui["animated"] = (True,)
|
||||
return IO.NodeOutput(images, ui=ui)
|
||||
|
||||
for batch_number, image in enumerate(images):
|
||||
encoded = _encode_image(image, file_format, bit_depth, colorspace)
|
||||
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import sys
|
||||
import torch
|
||||
from unittest.mock import patch, MagicMock
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
# Mock nodes module to prevent CUDA initialization during import
|
||||
mock_nodes = MagicMock()
|
||||
@@ -8,8 +9,19 @@ mock_nodes.MAX_RESOLUTION = 16384
|
||||
# Mock server module for PromptServer
|
||||
mock_server = MagicMock()
|
||||
|
||||
with patch.dict('sys.modules', {'nodes': mock_nodes, 'server': mock_server}):
|
||||
from comfy_extras.nodes_images import ImageStitch
|
||||
previous_nodes = sys.modules.get("nodes")
|
||||
previous_server = sys.modules.get("server")
|
||||
sys.modules["nodes"] = mock_nodes
|
||||
sys.modules["server"] = mock_server
|
||||
from comfy_extras.nodes_images import ImageStitch
|
||||
if previous_nodes is None:
|
||||
sys.modules.pop("nodes")
|
||||
else:
|
||||
sys.modules["nodes"] = previous_nodes
|
||||
if previous_server is None:
|
||||
sys.modules.pop("server")
|
||||
else:
|
||||
sys.modules["server"] = previous_server
|
||||
|
||||
|
||||
class TestImageStitch:
|
||||
@@ -240,4 +252,3 @@ class TestImageStitch:
|
||||
expected_image2_width = int(64 * (32/32)) # Resized to height 64
|
||||
expected_total_width = 48 + 8 + expected_image2_width
|
||||
assert result[0].shape[2] == expected_total_width
|
||||
|
||||
|
||||
Reference in New Issue
Block a user