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

add interpret_as

Browse Source
This commit is contained in:
Yousef Rafat 2026-04-29 21:23:28 +03:00
parent 5cca14c798
commit c77b36d98f
1 changed files with 23 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

40
comfy_extras/nodes_images.py
View File

@ -983,13 +983,13 @@ class SaveImageAdvanced(IO.ComfyNode):
[ [
IO.Combo.Input( IO.Combo.Input(
"bit_depth", "bit_depth",
options=["8-bit", "10-bit", "12-bit"], options=["8-bit", "10-bit"],
default="8-bit", default="8-bit",
advanced=True, advanced=True,
), ),
IO.Combo.Input( IO.Combo.Input(
"interpret_as", "interpret_as",
options=["sRGB", "Raw/Data"], options=["sRGB", "Linear", "Raw/Data"],
default="sRGB", default="sRGB",
advanced=True, advanced=True,
), ),
@ -1000,13 +1000,13 @@ class SaveImageAdvanced(IO.ComfyNode):
[ [
IO.Combo.Input( IO.Combo.Input(
"bit_depth", "bit_depth",
options=["16-bit", "32-bit"], options=["32-bit"],
default="16-bit", default="32-bit",
advanced=True, advanced=True,
), ),
IO.Combo.Input( IO.Combo.Input(
"interpret_as", "interpret_as",
options=["Linear", "Raw/Data"], options=["sRGB", "Linear", "Raw/Data"],
default="Linear", default="Linear",
advanced=True, advanced=True,
), ),
@ -1050,14 +1050,9 @@ class SaveImageAdvanced(IO.ComfyNode):
if bit_depth == "32-bit": if bit_depth == "32-bit":
img_np = img_tensor.cpu().numpy().astype(np.float32) img_np = img_tensor.cpu().numpy().astype(np.float32)
av_fmt = 'gbrapf32le' if has_alpha else 'gbrpf32le' av_fmt = 'gbrapf32le' if has_alpha else 'gbrpf32le'
elif bit_depth == "16-bit": elif bit_depth in ["10-bit", "12-bit", "16-bit"]:
if file_format == "exr": img_np = (img_tensor * 65535.0).clamp(0, 65535).to(torch.int32).cpu().numpy().astype(np.uint16)
# default pyav build doesn't come with a codec for float16 exr format av_fmt = 'rgba64le' if has_alpha else 'rgb48le'
img_np = img_tensor.cpu().numpy().astype(np.float32)
av_fmt = 'gbrapf32le' if has_alpha else 'gbrpf32le'
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: else:
img_np = (img_tensor * 255.0).clamp(0, 255).to(torch.int32).cpu().numpy().astype(np.uint8) img_np = (img_tensor * 255.0).clamp(0, 255).to(torch.int32).cpu().numpy().astype(np.uint8)
av_fmt = 'rgba' if has_alpha else 'rgb24' av_fmt = 'rgba' if has_alpha else 'rgb24'
@ -1079,15 +1074,26 @@ class SaveImageAdvanced(IO.ComfyNode):
stream.time_base = Fraction(1, 1) stream.time_base = Fraction(1, 1)
if bit_depth in ["16-bit", "32-bit"]: if bit_depth == "12-bit":
stream.pix_fmt = 'yuv420p12le'
elif bit_depth in ["10-bit", "16-bit", "32-bit"]:
stream.pix_fmt = 'yuv420p10le' stream.pix_fmt = 'yuv420p10le'
else: else:
stream.pix_fmt = 'yuv420p' stream.pix_fmt = 'yuv420p'
stream.codec_context.color_range = 2 stream.codec_context.color_range = 2
stream.codec_context.colorspace = 1 if interpret_as == "Raw/Data": # 2 == unspecified
stream.codec_context.color_primaries = 1 stream.codec_context.colorspace = 2
stream.codec_context.color_trc = 1 stream.codec_context.color_primaries = 2
stream.codec_context.color_trc = 2
elif interpret_as == "Linear":
stream.codec_context.colorspace = 1
stream.codec_context.color_primaries = 1
stream.codec_context.color_trc = 8
else: # sRGB
stream.codec_context.colorspace = 1
stream.codec_context.color_primaries = 1
stream.codec_context.color_trc = 13
stream.options = { stream.options = {
'preset': '10', 'preset': '10',