From 537e3039dca48f247413d82703cac6f2a1a3f76a Mon Sep 17 00:00:00 2001
From: Yousef Rafat <81116377+yousef-rafat@users.noreply.github.com>
Date: Fri, 3 Jul 2026 18:01:57 +0300
Subject: [PATCH] .

---
 comfy_extras/nodes_images.py | 25 +++++++++++++++++++++----
 1 file changed, 21 insertions(+), 4 deletions(-)

diff --git a/comfy_extras/nodes_images.py b/comfy_extras/nodes_images.py
index fe1937ba5..bd3e76adc 100644
--- a/comfy_extras/nodes_images.py
+++ b/comfy_extras/nodes_images.py
@@ -1101,11 +1101,28 @@ def _encode_image(
     For PNG, colorspace selection does not modify pixels — PNG is delivered
     sRGB-encoded and there is no PNG path for wide-gamut HDR in this node.
     """
+
+    if img_tensor.ndim == 2: # 2 channel grayscales
+        img_tensor = img_tensor.unsqueeze(-1)
+
     height, width, num_channels = img_tensor.shape
     has_alpha = num_channels == 4
 
     spec = _FORMAT_SPECS[(file_format, bit_depth, has_alpha)]
 
+    frame_fmt = spec["frame_fmt"]
+    stream_fmt = spec["stream_fmt"]
+
+    if num_channels == 1:
+        if spec["dtype"] == np.float32:
+            frame_fmt = stream_fmt = "grayf32le"
+        elif spec["dtype"] == np.uint16:
+            frame_fmt = "gray16le"
+            # Ensure big-endian for PNG
+            stream_fmt = "gray16be" if "be" in stream_fmt else "gray16le"
+        else:
+            frame_fmt = stream_fmt = "gray"
+
     if spec["dtype"] == np.float32:
         # EXR path: preserve full range, no clamp.
         if colorspace == "sRGB":
@@ -1124,12 +1141,12 @@ def _encode_image(
     codec = av.CodecContext.create(file_format, "w")
     codec.width = width
     codec.height = height
-    codec.pix_fmt = spec["stream_fmt"]
+    codec.pix_fmt = stream_fmt
     codec.time_base = Fraction(1, 1)
 
-    frame = av.VideoFrame.from_ndarray(img_np, format=spec["frame_fmt"])
-    if spec["frame_fmt"] != spec["stream_fmt"]:
-        frame = frame.reformat(format=spec["stream_fmt"])
+    frame = av.VideoFrame.from_ndarray(img_np, format=frame_fmt)
+    if frame_fmt != stream_fmt:
+        frame = frame.reformat(format=stream_fmt)
     frame.pts = 0
     frame.time_base = codec.time_base