wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-05-29 02:17:52 +08:00
Code Issues Actions 30 Packages Projects Releases Wiki Activity

correct avif code

Browse Source
This commit is contained in:
Yousef Rafat 2026-04-27 21:52:44 +03:00
parent ebb9acf3cf
commit 80e421b626
1 changed files with 37 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

44
comfy_extras/nodes_save_advanced.py
View File

@ -12,6 +12,7 @@ from comfy_api.latest import IO
from typing_extensions import override from typing_extensions import override
from comfy_api.latest import ComfyExtension from comfy_api.latest import ComfyExtension
from comfy.cli_args import args from comfy.cli_args import args
from fractions import Fraction
def create_png_chunk(chunk_type: bytes, data: bytes) -> bytes: def create_png_chunk(chunk_type: bytes, data: bytes) -> bytes:
"""Creates a valid PNG chunk with Length, Type, Data, and CRC32.""" """Creates a valid PNG chunk with Length, Type, Data, and CRC32."""
@ -189,10 +190,31 @@ class SaveImageAdvanced(IO.ComfyNode):
if file_format == "exr": if file_format == "exr":
stream = container.add_stream('exr', rate=1) stream = container.add_stream('exr', rate=1)
stream.pix_fmt = av_fmt stream.pix_fmt = av_fmt
elif file_format == "avif": elif file_format == "avif":
stream = container.add_stream('av1', rate=1) try:
# YUV color spaces stream = container.add_stream('libsvtav1', rate=1)
stream.pix_fmt = 'yuv444p12le' if bit_depth in["16-bit", "32-bit"] else 'yuv444p' except Exception:
stream = container.add_stream('av1', rate=1)
stream.time_base = Fraction(1, 1)
if bit_depth in ["16-bit", "32-bit"]:
stream.pix_fmt = 'yuv420p10le'
else:
stream.pix_fmt = 'yuv420p'
stream.codec_context.color_range = 2
stream.codec_context.colorspace = 1
stream.codec_context.color_primaries = 1
stream.codec_context.color_trc = 1
stream.options = {
'preset': '10',
'svtav1-params': 'rc=0:qp=20:color-range=1:color-matrix=1:enable-overlays=1',
'g': '1'
}
elif file_format == "png": elif file_format == "png":
stream = container.add_stream('png', rate=1) stream = container.add_stream('png', rate=1)
stream.pix_fmt = av_fmt stream.pix_fmt = av_fmt
@ -208,8 +230,6 @@ class SaveImageAdvanced(IO.ComfyNode):
try: try:
frame = av.VideoFrame.from_ndarray(img_np, format=av_fmt) frame = av.VideoFrame.from_ndarray(img_np, format=av_fmt)
except ValueError: except ValueError:
# FFMPEG Float32 Fallback: not all ffmpeg versions are able to handle float32 format for images
# float16 fallback conversion
logging.warning("[WARNING] Current FFMPEG Binary can't save float32 images. Fallbacking to float16") logging.warning("[WARNING] Current FFMPEG Binary can't save float32 images. Fallbacking to float16")
img_np = (img_tensor * 65535.0).clamp(0, 65535).to(torch.int32).cpu().numpy().astype(np.uint16) img_np = (img_tensor * 65535.0).clamp(0, 65535).to(torch.int32).cpu().numpy().astype(np.uint16)
av_fmt = 'rgba64le' if has_alpha else 'rgb48le' av_fmt = 'rgba64le' if has_alpha else 'rgb48le'
@ -217,6 +237,18 @@ class SaveImageAdvanced(IO.ComfyNode):
if file_format == "exr" or file_format == "png": if file_format == "exr" or file_format == "png":
stream.pix_fmt = av_fmt stream.pix_fmt = av_fmt
# reformat for avif
if file_format == "avif":
frame = frame.reformat(
format=stream.pix_fmt,
src_colorspace=1, dst_colorspace=1,
src_color_range=2, dst_color_range=2
)
frame.pts = 0
frame.time_base = stream.time_base
frame.color_range = 2
frame.colorspace = 1
for packet in stream.encode(frame): for packet in stream.encode(frame):
container.mux(packet) container.mux(packet)
for packet in stream.encode(): for packet in stream.encode():
@ -246,8 +278,6 @@ class SaveImageAdvanced(IO.ComfyNode):
}) })
counter += 1 counter += 1
return IO.NodeOutput(ui={"images": results})
# Rec.709 to Rec.2020 Gamut Conversion Matrix # Rec.709 to Rec.2020 Gamut Conversion Matrix
M_709_to_2020 = torch.tensor([[0.6274, 0.3293, 0.0433],[0.0691, 0.9195, 0.0114],[0.0164, 0.0880, 0.8956] M_709_to_2020 = torch.tensor([[0.6274, 0.3293, 0.0433],[0.0691, 0.9195, 0.0114],[0.0164, 0.0880, 0.8956]
]) ])