wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-05-18 13:07:28 +08:00
Code Issues Actions 7 Packages Projects Releases Wiki Activity

png, exr, avif

Browse Source
This commit is contained in:
Yousef Rafat 2026-04-28 21:40:13 +03:00
parent 80e421b626
commit 29653e7fcc
1 changed files with 28 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

50
comfy_extras/nodes_save_advanced.py
View File

@ -167,16 +167,14 @@ class SaveImageAdvanced(IO.ComfyNode):
bit_depth = "16-bit"
if bit_depth == "32-bit":
img_np = img_tensor.cpu().numpy()
# rgba128le handles 4x32f, gbrpf32le handles 3x32f planar
av_fmt = 'rgba128le' if has_alpha else 'gbrpf32le'
img_np = img_tensor.cpu().numpy().astype(np.float32)
av_fmt = 'gbrapf32le' if has_alpha else 'gbrpf32le'
elif bit_depth == "16-bit":
img_np = (img_tensor * 65535.0).clamp(0, 65535).to(torch.int32).cpu().numpy().astype(np.uint16)
if file_format == "png":
# png requires Big-Endian (be) for 16-bit
av_fmt = 'rgba64be' if has_alpha else 'rgb48be'
img_np = img_np.byteswap().view(img_np.dtype.newbyteorder('>'))
if file_format == "exr":
img_np = img_tensor.cpu().numpy().astype(np.float16)
av_fmt = 'gbrapf16le' if has_alpha else 'gbrpf16le'
else:
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'
else:
img_np = (img_tensor * 255.0).clamp(0, 255).to(torch.int32).cpu().numpy().astype(np.uint8)
@ -217,37 +215,43 @@ class SaveImageAdvanced(IO.ComfyNode):
elif file_format == "png":
stream = container.add_stream('png', rate=1)
stream.pix_fmt = av_fmt
if bit_depth == "16-bit":
stream.pix_fmt = 'rgba64be' if has_alpha else 'rgb48be'
else:
stream.pix_fmt = av_fmt
stream.width = width
stream.height = height
stream.time_base = Fraction(1, 1)
# planar: all red, all blue, all green instead of r, g, b, r, g, b
is_planar = av_fmt.startswith('gbrp') or 'p' in av_fmt.split('rgba')[-1]
if is_planar:
if av_fmt.startswith('gbrp'):
img_np = img_np[:, :, [1, 2, 0, 3]] if has_alpha else img_np[:, :, [1, 2, 0]]
img_np = img_np.transpose(2, 0, 1)
try:
frame = av.VideoFrame.from_ndarray(img_np, format=av_fmt)
except ValueError:
logging.warning("[WARNING] Current FFMPEG Binary can't save float32 images. Fallbacking to float16")
logging.warning("[WARNING] Current FFMPEG Binary can't save natively. Fallbacking.")
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'
frame = av.VideoFrame.from_ndarray(img_np, format=av_fmt)
if file_format == "exr" or file_format == "png":
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
)
# reformat for both avif and exr to ensure correct internal conversion
if file_format in ["avif", "exr"] or (file_format == "png" and bit_depth == "16-bit"):
reformat_kwargs = {"format": stream.pix_fmt}
if file_format == "avif":
reformat_kwargs.update({
"src_colorspace": 1, "dst_colorspace": 1,
"src_color_range": 2, "dst_color_range": 2
})
frame = frame.reformat(**reformat_kwargs)
frame.pts = 0
frame.time_base = stream.time_base
frame.color_range = 2
frame.colorspace = 1
if file_format == "avif":
frame.color_range = 2
frame.colorspace = 1
for packet in stream.encode(frame):
container.mux(packet)
@ -278,6 +282,8 @@ class SaveImageAdvanced(IO.ComfyNode):
})
counter += 1
return IO.NodeOutput(ui={"images": results})
# 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]
])