Add 10-bit video support (#14452) · code/app-python-comfyui@fe54b5e · GitHub
Skip to content

Commit fe54b5e

Browse files
authored
Add 10-bit video support (Comfy-Org#14452)
Create Video gets a bit_depth option (8-bit/10-bit); the selected depth is carried by the video and applied when it gets encoded. Save Video and Video Slice now keep the source bit depth instead of always quantizing to 8-bit, so 10-bit videos stay 10-bit. 10-bit uses h264 with the yuv420p10le pixel format,so there's no new codec or container. Signed-off-by: bigcat88 <bigcat88@icloud.com>
1 parent 7277d99 commit fe54b5e

4 files changed

Lines changed: 169 additions & 12 deletions

File tree

‎comfy_api/latest/_input/video_types.py‎

Lines changed: 12 additions & 1 deletion

‎comfy_api/latest/_input_impl/video_types.py‎

Lines changed: 43 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,12 @@ def get_open_write_kwargs(
5252
return open_kwargs
5353

5454

55+
def video_stream_bit_depth(stream) -> int:
56+
if stream is None or stream.format is None or not stream.format.components:
57+
return 8
58+
return max(component.bits for component in stream.format.components)
59+
60+
5561
class VideoFromFile(VideoInput):
5662
"""
5763
Class representing video input from a file.
@@ -97,6 +103,13 @@ def get_dimensions(self) -> tuple[int, int]:
97103
return stream.width, stream.height
98104
raise ValueError(f"No video stream found in file '{self.__file}'")
99105

106+
def get_bit_depth(self) -> int:
107+
if isinstance(self.__file, io.BytesIO):
108+
self.__file.seek(0) # Reset the BytesIO object to the beginning
109+
with av.open(self.__file, mode="r") as container:
110+
video_stream = container.streams.video[0] if len(container.streams.video) > 0 else None
111+
return video_stream_bit_depth(video_stream)
112+
100113
def get_duration(self) -> float:
101114
"""
102115
Returns the duration of the video in seconds.
@@ -377,25 +390,32 @@ def save_to(
377390
format: VideoContainer = VideoContainer.AUTO,
378391
codec: VideoCodec = VideoCodec.AUTO,
379392
metadata: Optional[dict] = None,
393+
bit_depth: int | None = None,
380394
):
381395
if isinstance(self.__file, io.BytesIO):
382396
self.__file.seek(0) # Reset the BytesIO object to the beginning
383397
with av.open(self.__file, mode='r') as container:
384398
container_format = container.format.name
385-
video_encoding = container.streams.video[0].codec.name if len(container.streams.video) > 0 else None
399+
video_stream = container.streams.video[0] if len(container.streams.video) > 0 else None
400+
video_encoding = video_stream.codec.name if video_stream is not None else None
401+
source_bit_depth = video_stream_bit_depth(video_stream)
386402
reuse_streams = True
387403
if format != VideoContainer.AUTO and format not in container_format.split(","):
388404
reuse_streams = False
389405
if codec != VideoCodec.AUTO and codec != video_encoding and video_encoding is not None:
390406
reuse_streams = False
407+
if bit_depth is not None and video_encoding is not None and bit_depth != source_bit_depth:
408+
reuse_streams = False
391409
if self.__start_time or self.__duration:
392410
reuse_streams = False
393411

394412
if not reuse_streams:
413+
if bit_depth is None:
414+
bit_depth = source_bit_depth
395415
components = self.get_components_internal(container)
396416
video = VideoFromComponents(components)
397417
return video.save_to(
398-
path, format=format, codec=codec, metadata=metadata
418+
path, format=format, codec=codec, metadata=metadata, bit_depth=bit_depth,
399419
)
400420

401421
streams = container.streams
@@ -451,8 +471,10 @@ class VideoFromComponents(VideoInput):
451471
Class representing video input from tensors.
452472
"""
453473

454-
def __init__(self, components: VideoComponents):
474+
def __init__(self, components: VideoComponents, bit_depth: int = 8):
455475
self.__components = components
476+
# Tensor components have no inherent bit depth; this is the depth used when encoding.
477+
self.__bit_depth = bit_depth
456478

457479
def get_components(self) -> VideoComponents:
458480
return VideoComponents(
@@ -461,18 +483,26 @@ def get_components(self) -> VideoComponents:
461483
frame_rate=self.__components.frame_rate,
462484
)
463485

486+
def get_bit_depth(self) -> int:
487+
return self.__bit_depth
488+
464489
def save_to(
465490
self,
466491
path: str,
467492
format: VideoContainer = VideoContainer.AUTO,
468493
codec: VideoCodec = VideoCodec.AUTO,
469494
metadata: Optional[dict] = None,
495+
bit_depth: int | None = None,
470496
):
471497
"""Save the video to a file path or BytesIO buffer."""
472498
if format != VideoContainer.AUTO and format != VideoContainer.MP4:
473499
raise ValueError("Only MP4 format is supported for now")
474500
if codec != VideoCodec.AUTO and codec != VideoCodec.H264:
475501
raise ValueError("Only H264 codec is supported for now")
502+
# None means "use the depth this video was created with" (CreateVideo's choice).
503+
if bit_depth is None:
504+
bit_depth = self.__bit_depth
505+
is_10bit = bit_depth >= 10
476506
extra_kwargs = {}
477507
if isinstance(format, VideoContainer) and format != VideoContainer.AUTO:
478508
extra_kwargs["format"] = format.value
@@ -488,10 +518,11 @@ def save_to(
488518

489519
frame_rate = Fraction(round(self.__components.frame_rate * 1000), 1000)
490520
# Create a video stream
521+
pix_fmt = "yuv420p10le" if is_10bit else "yuv420p"
491522
video_stream = output.add_stream('h264', rate=frame_rate)
492523
video_stream.width = self.__components.images.shape[2]
493524
video_stream.height = self.__components.images.shape[1]
494-
video_stream.pix_fmt = 'yuv420p'
525+
video_stream.pix_fmt = pix_fmt
495526

496527
# Create an audio stream
497528
audio_sample_rate = 1
@@ -505,9 +536,14 @@ def save_to(
505536

506537
# Encode video
507538
for i, frame in enumerate(self.__components.images):
508-
img = (frame * 255).clamp(0, 255).byte().cpu().numpy() # shape: (H, W, 3)
509-
frame = av.VideoFrame.from_ndarray(img, format='rgb24')
510-
frame = frame.reformat(format='yuv420p') # Convert to YUV420P as required by h264
539+
if is_10bit:
540+
# 16-bit RGB keeps float precision through the conversion to 10-bit YUV.
541+
img = (frame.float() * 65535).clamp(0, 65535).cpu().numpy().astype(np.uint16) # shape: (H, W, 3)
542+
frame = av.VideoFrame.from_ndarray(img, format="rgb48le")
543+
else:
544+
img = (frame * 255).clamp(0, 255).byte().cpu().numpy() # shape: (H, W, 3)
545+
frame = av.VideoFrame.from_ndarray(img, format='rgb24')
546+
frame = frame.reformat(format=pix_fmt)
511547
packet = video_stream.encode(frame)
512548
output.mux(packet)
513549

‎comfy_extras/nodes_video.py‎

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -134,16 +134,32 @@ def define_schema(cls):
134134
io.Image.Input("images", tooltip="The images to create a video from."),
135135
io.Float.Input("fps", default=30.0, min=1.0, max=120.0, step=1.0),
136136
io.Audio.Input("audio", optional=True, tooltip="The audio to add to the video."),
137+
io.Int.Input(
138+
"bit_depth",
139+
min=8,
140+
max=10,
141+
default=8,
142+
step=2,
143+
tooltip="Bit depth of the created video. 10-bit keeps smoother gradients with less"
144+
" banding, but some players and downstream nodes may not support it.",
145+
optional=True,
146+
display_mode=io.NumberDisplay.number,
147+
),
137148
],
138149
outputs=[
139150
io.Video.Output(),
140151
],
141152
)
142153

143154
@classmethod
144-
def execute(cls, images: Input.Image, fps: float, audio: Optional[Input.Audio] = None) -> io.NodeOutput:
155+
def execute(
156+
cls, images: Input.Image, fps: float, audio: Optional[Input.Audio] = None, bit_depth: int = 8,
157+
) -> io.NodeOutput:
145158
return io.NodeOutput(
146-
InputImpl.VideoFromComponents(Types.VideoComponents(images=images, audio=audio, frame_rate=Fraction(fps)))
159+
InputImpl.VideoFromComponents(
160+
Types.VideoComponents(images=images, audio=audio, frame_rate=Fraction(fps)),
161+
bit_depth=bit_depth,
162+
)
147163
)
148164

149165
class GetVideoComponents(io.ComfyNode):
@@ -154,21 +170,22 @@ def define_schema(cls):
154170
search_aliases=["extract frames", "split video", "video to images", "demux"],
155171
display_name="Get Video Components",
156172
category="video",
157-
description="Extracts all components from a video: frames, audio, and framerate.",
173+
description="Extracts all components from a video: frames, audio, framerate, and bit depth.",
158174
inputs=[
159175
io.Video.Input("video", tooltip="The video to extract components from."),
160176
],
161177
outputs=[
162178
io.Image.Output(display_name="images"),
163179
io.Audio.Output(display_name="audio"),
164180
io.Float.Output(display_name="fps"),
181+
io.Int.Output(display_name="bit_depth"),
165182
],
166183
)
167184

168185
@classmethod
169186
def execute(cls, video: Input.Video) -> io.NodeOutput:
170187
components = video.get_components()
171-
return io.NodeOutput(components.images, components.audio, float(components.frame_rate))
188+
return io.NodeOutput(components.images, components.audio, float(components.frame_rate), video.get_bit_depth())
172189

173190

174191
class LoadVideo(io.ComfyNode):
Lines changed: 93 additions & 0 deletions

0 commit comments

Comments
 (0)