From 1d595d4a703515664157849daeaee88d7ca47a1c Mon Sep 17 00:00:00 2001 From: bigcat88 Date: Tue, 7 Jul 2026 00:06:14 +0300 Subject: [PATCH 1/3] fix(Video): stream the video transcode instead of buffering every frame in RAM Signed-off-by: bigcat88 --- comfy_api/latest/_input_impl/video_types.py | 366 ++++++++++++++--- tests-unit/comfy_api_test/video_types_test.py | 386 +++++++++++++++++- 2 files changed, 700 insertions(+), 52 deletions(-) diff --git a/comfy_api/latest/_input_impl/video_types.py b/comfy_api/latest/_input_impl/video_types.py index bc95a5b99..b9511e593 100644 --- a/comfy_api/latest/_input_impl/video_types.py +++ b/comfy_api/latest/_input_impl/video_types.py @@ -1,5 +1,6 @@ from av.container import InputContainer from av.subtitles.stream import SubtitleStream +from av.video.reformatter import ColorRange from fractions import Fraction from typing import Optional from .._input import AudioInput, VideoInput @@ -9,6 +10,7 @@ import itertools import json import numpy as np import math +import os import torch from .._util import VideoContainer, VideoCodec, VideoComponents import logging @@ -58,6 +60,57 @@ def video_stream_bit_depth(stream) -> int: return max(component.bits for component in stream.format.components) +def last_decodable_audio_stream(container: InputContainer): + """Streams FFmpeg has no decoder for have no codec context, and decoding their + packets crashes the process (e.g. APAC spatial-audio track in iPhone).""" + stream = next( + (s for s in reversed(container.streams.audio) if s.codec_context is not None), + None, + ) + if stream is None and len(container.streams.audio): + logging.warning("No decodable audio stream found in video; ignoring audio.") + return stream + + +def probe_audio_params(container: InputContainer, audio_stream, max_packets: int = 200): + """Containers probed only up to a window (mpegts) leave audio codec parameters unset when + audio starts beyond it; learn them by decoding ahead. The caller must seek back afterwards. + Returns (sample_rate, channels), zeros when the stream never yields a decodable frame.""" + for i, packet in enumerate(container.demux(audio_stream)): + try: + frames = packet.decode() + except av.error.FFmpegError: + return 0, 0 + if frames: + return frames[0].sample_rate, frames[0].layout.nb_channels + if i >= max_packets: + break + return 0, 0 + + +def write_output_metadata(container: InputContainer, output, metadata: dict | None): + """Copy the source container's metadata, then overlay the caller's tags.""" + for key, value in container.metadata.items(): + if metadata is None or key not in metadata: + output.metadata[key] = value + if metadata is not None: + for key, value in metadata.items(): + output.metadata[key] = value if isinstance(value, str) else json.dumps(value) + + +def mp4_output_open_kwargs(path: str | io.BytesIO, format: VideoContainer, codec: VideoCodec) -> dict: + if format != VideoContainer.AUTO and format != VideoContainer.MP4: + raise ValueError("Only MP4 format is supported for now") + if codec != VideoCodec.AUTO and codec != VideoCodec.H264: + raise ValueError("Only H264 codec is supported for now") + open_kwargs = {"mode": "w", "options": {"movflags": "use_metadata_tags"}} + if isinstance(format, VideoContainer) and format != VideoContainer.AUTO: + open_kwargs["format"] = format.value + elif isinstance(path, io.BytesIO): + open_kwargs["format"] = "mp4" # no file extension to infer the format from + return open_kwargs + + class VideoFromFile(VideoInput): """ Class representing video input from a file. @@ -192,13 +245,10 @@ class VideoFromFile(VideoInput): return estimated_frames # 3. Last resort: decode frames and count them (streaming) - if self.__start_time < 0: - start_time = max(self._get_raw_duration() + self.__start_time, 0) - else: - start_time = self.__start_time + start_time, duration = self.get_active_trim_window() frame_count = 1 start_pts = int(start_time / video_stream.time_base) - end_pts = int((start_time + self.__duration) / video_stream.time_base) + end_pts = int((start_time + duration) / video_stream.time_base) container.seek(start_pts, stream=video_stream) frame_iterator = ( container.decode(video_stream) @@ -253,17 +303,14 @@ class VideoFromFile(VideoInput): def get_components_internal(self, container: InputContainer) -> VideoComponents: video_stream = self._get_first_video_stream(container) - if self.__start_time < 0: - start_time = max(self._get_raw_duration() + self.__start_time, 0) - else: - start_time = self.__start_time + start_time, duration = self.get_active_trim_window() # Get video frames frames = [] audio_frames = [] alphas = None start_pts = int(start_time / video_stream.time_base) - end_pts = int((start_time + self.__duration) / video_stream.time_base) + end_pts = int((start_time + duration) / video_stream.time_base) if start_pts != 0: container.seek(start_pts, stream=video_stream) @@ -281,18 +328,11 @@ class VideoFromFile(VideoInput): video_done = False audio_done = True - # Use the last decodable audio stream. Streams FFmpeg has no decoder for have no codec context, - # and decoding their packets crashes the process. (e.g. APAC spatial-audio track in iPhone) - audio_stream = next( - (s for s in reversed(container.streams.audio) if s.codec_context is not None), - None, - ) + audio_stream = last_decodable_audio_stream(container) if audio_stream is not None: streams += [audio_stream] resampler = av.audio.resampler.AudioResampler(format='fltp') audio_done = False - elif len(container.streams.audio): - logging.warning("No decodable audio stream found in video; ignoring audio.") for packet in container.demux(*streams): if video_done and audio_done: @@ -305,7 +345,7 @@ class VideoFromFile(VideoInput): for frame in packet.decode(): if frame.pts < start_pts: continue - if self.__duration and frame.pts >= end_pts: + if duration and frame.pts >= end_pts: video_done = True break @@ -372,7 +412,7 @@ class VideoFromFile(VideoInput): map(resampler.resample, packet.decode()) ) for frame in aframes: - if self.__duration and frame.time > start_time + self.__duration: + if duration and frame.time > start_time + duration: audio_done = True break @@ -394,8 +434,8 @@ class VideoFromFile(VideoInput): if len(audio_frames) > 0: audio_data = np.concatenate(audio_frames, axis=1) # shape: (channels, total_samples) - if self.__duration: - audio_data = audio_data[..., :int(self.__duration * audio_stream.sample_rate)] + if duration: + audio_data = audio_data[..., :int(duration * audio_stream.sample_rate)] audio_tensor = torch.from_numpy(audio_data).unsqueeze(0) # shape: (1, channels, total_samples) audio = AudioInput({ @@ -441,28 +481,14 @@ class VideoFromFile(VideoInput): if not reuse_streams: if bit_depth is None: bit_depth = source_bit_depth - components = self.get_components_internal(container) - video = VideoFromComponents(components) - return video.save_to( - path, format=format, codec=codec, metadata=metadata, bit_depth=bit_depth, - ) + return self._save_transcoded(container, path, format=format, codec=codec, metadata=metadata, bit_depth=bit_depth) streams = container.streams open_kwargs = get_open_write_kwargs(path, container_format, format) with av.open(path, **open_kwargs) as output_container: - # Copy over the original metadata - for key, value in container.metadata.items(): - if metadata is None or key not in metadata: - output_container.metadata[key] = value - - # Add our new metadata - if metadata is not None: - for key, value in metadata.items(): - if isinstance(value, str): - output_container.metadata[key] = value - else: - output_container.metadata[key] = json.dumps(value) + # Add metadata before writing any streams + write_output_metadata(container, output_container, metadata) # Add streams to the new container. Streams with no codec context cannot be used as an output template. stream_map = {} @@ -480,6 +506,254 @@ class VideoFromFile(VideoInput): packet.stream = stream_map[packet.stream] output_container.mux(packet) + def _save_transcoded( + self, + container: InputContainer, + path: str | io.BytesIO, + format: VideoContainer, + codec: VideoCodec, + metadata: dict | None, + bit_depth: int, + ): + """Re-encode to H.264/AAC one frame at a time; peak memory does not scale with video length.""" + open_kwargs = mp4_output_open_kwargs(path, format, codec) + video_stream = self._get_first_video_stream(container) + start_time, duration = self.get_active_trim_window() + start_pts = int(start_time / video_stream.time_base) + end_pts = int((start_time + duration) / video_stream.time_base) if duration else None + if start_pts != 0: + container.seek(start_pts, stream=video_stream) + + audio_stream = last_decodable_audio_stream(container) + pix_fmt = "yuv420p10le" if bit_depth >= 10 else "yuv420p" + rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1) + + resampler = None + sample_rate = 0 + audio_time_base = None + duration_cap = None + if audio_stream is not None: + sample_rate = audio_stream.codec_context.sample_rate + channels = audio_stream.codec_context.channels + if not sample_rate: + sample_rate, channels = probe_audio_params(container, audio_stream) + container.seek(start_pts, stream=video_stream) + if sample_rate: + audio_stream.codec_context.flush_buffers() + else: + logging.warning("Audio stream parameters could not be determined; ignoring audio.") + audio_stream = None + if audio_stream is not None: + audio_time_base = Fraction(1, sample_rate) + layout = {1: "mono", 2: "stereo", 6: "5.1"}.get(channels, "stereo") + resampler = av.audio.resampler.AudioResampler(format="fltp", layout=layout, rate=sample_rate) + if duration: + duration_cap = math.ceil(duration * sample_rate) + + streams = [video_stream] if audio_stream is None else [video_stream, audio_stream] + pts_step = max(1, int(round((1 / rate) / video_stream.time_base))) + video_done = False + audio_done = audio_stream is None + video_pts_offset = None + last_video_pts = None + last_video_end = None + # rebased pts -> true display duration: the mp4 muxer pads the last sample with 1/rate otherwise + video_frame_durations = {} + source_size = None + rotation_k = 0 + rotation_filter = None + audio_started = False + samples_written = 0 + pending_audio = [] + # The output opens lazily on the first kept frame: it decides the geometry (90/270 rotation swaps dims), + # and never seeking back keeps webm/mkv leading audio intact. + output = None + out_video = None + out_audio = None + + def audio_frame_from_ndarray(nd_planar): + frame = av.AudioFrame.from_ndarray(np.ascontiguousarray(nd_planar), format="fltp", layout=layout) + frame.sample_rate = sample_rate + return frame + + def drain_audio(final=False): + # Audio may cover the pts span of the video written so far, capped by the requested duration + nonlocal samples_written, audio_done + if last_video_end is None: + cap = 0 + else: + cap = math.ceil(last_video_end * video_stream.time_base * sample_rate) + if duration_cap is not None: + cap = min(cap, duration_cap) + while pending_audio and not audio_done: + frame = pending_audio[0] + if samples_written + frame.samples <= cap: + frame.pts = samples_written + frame.time_base = audio_time_base + output.mux(out_audio.encode(frame)) + samples_written += frame.samples + pending_audio.pop(0) + continue + if final: + keep = frame.to_ndarray()[..., :cap - samples_written] + if keep.shape[-1] > 0: + tail = audio_frame_from_ndarray(keep) + tail.pts = samples_written + tail.time_base = audio_time_base + output.mux(out_audio.encode(tail)) + samples_written += keep.shape[-1] + pending_audio.clear() + break + if duration_cap is not None and samples_written >= duration_cap: + audio_done = True + return cap + + try: + for packet in container.demux(*streams): + if video_done and audio_done: + break + + if packet.stream == video_stream and not video_done: + try: + frames = packet.decode() + except av.error.InvalidDataError: + logging.info("pyav decode error") + continue + for frame in frames: + if frame.pts is not None and frame.pts < start_pts: + continue + if end_pts is not None and frame.pts is not None and frame.pts >= end_pts: + video_done = True + if last_video_pts is not None: + # the source continues past the window: hold the last kept frame to the window end + last_video_end = max(last_video_end, end_pts - video_pts_offset) + break + # the source's true display duration of this frame; average_rate is not a + # frame duration (sparse/VFR sources), so it is only the fallback + frame_duration = frame.duration if frame.duration else pts_step + if end_pts is not None and frame.pts is not None: + frame_duration = min(frame_duration, end_pts - frame.pts) + if output is None: + rotation_k = int(round(frame.rotation // 90)) % 4 if frame.rotation else 0 + if rotation_k % 2: + out_width, out_height = frame.height, frame.width + else: + out_width, out_height = frame.width, frame.height + if out_width % 2 or out_height % 2: + raise ValueError(f"H.264 output requires even dimensions, got {out_width}x{out_height}") + source_size = (frame.width, frame.height) + output = av.open(path, **open_kwargs) + # Add metadata before writing any streams + write_output_metadata(container, output, metadata) + out_video = output.add_stream("h264", rate=rate) + # no B-frames: reordering makes mp4 sample durations follow decode order, + # so irregular-VFR spans and trim windows land wrong + out_video.codec_context.max_b_frames = 0 + out_video.width = out_width + out_video.height = out_height + out_video.pix_fmt = pix_fmt + # source pts pass through (rebased to 0), so variable frame rate survives + out_video.codec_context.time_base = video_stream.time_base + if audio_stream is not None: + out_audio = output.add_stream("aac", rate=sample_rate, layout=layout) + if (frame.width, frame.height) != source_size: + # encoding would silently rescale the new geometry into the old one + raise ValueError( + f"Video resolution changes mid-stream " + f"({source_size[0]}x{source_size[1]} -> {frame.width}x{frame.height}); cannot transcode" + ) + if rotation_k: + if rotation_filter is None: + g = av.filter.Graph() + g_src = g.add_buffer(width=frame.width, height=frame.height, + format=frame.format.name, time_base=video_stream.time_base) + tail = g_src + for filter_name, filter_args in {1: [("transpose", "cclock")], + 2: [("hflip", None), ("vflip", None)], + 3: [("transpose", "clock")]}[rotation_k]: + step = g.add(filter_name, filter_args) + tail.link_to(step) + tail = step + g_sink = g.add("buffersink") + tail.link_to(g_sink) + g.configure() + rotation_filter = (g_src, g_sink) + rotation_filter[0].push(frame) + frame = rotation_filter[1].pull() + if frame.color_range == ColorRange.JPEG: + # compress full-range sources (yuvj/MJPEG) to limited range + frame = frame.reformat(format=pix_fmt, src_color_range="JPEG", dst_color_range="MPEG") + else: + frame = frame.reformat(format=pix_fmt) + if frame.pts is not None: + if video_pts_offset is None: + video_pts_offset = frame.pts + frame.pts -= video_pts_offset + if frame.pts is None or (last_video_pts is not None and frame.pts <= last_video_pts): + # broken sources emit missing/backward timestamps mid-stream, which the + # muxer rejects; nudge them forward by one nominal frame interval + frame.pts = 0 if last_video_pts is None else last_video_pts + pts_step + last_video_pts = frame.pts + last_video_end = frame.pts + frame_duration + video_frame_durations[frame.pts] = frame_duration + # the decoded pict_type would force x264's frame types (intra-only + # sources like MJPEG/ProRes would come out all-keyframe) + frame.pict_type = 0 + for out_packet in out_video.encode(frame): + out_packet.duration = video_frame_durations.pop(out_packet.pts, 0) + output.mux(out_packet) + drain_audio() + + elif packet.stream == audio_stream and not audio_done: + for resampled in itertools.chain.from_iterable(map(resampler.resample, packet.decode())): + if not audio_started: + if resampled.pts is None: + frame_start = 0.0 + else: + # passthrough frames keep the source stream's time base + tb = resampled.time_base if resampled.time_base else audio_time_base + frame_start = float(resampled.pts * tb) + to_skip = max(0, int((start_time - frame_start) * sample_rate)) + if to_skip >= resampled.samples: + continue + audio_started = True + if to_skip: + pending_audio.append(audio_frame_from_ndarray(resampled.to_ndarray()[..., to_skip:])) + continue + pending_audio.append(resampled) + if video_done: + # the video window is complete so the cap is final, but containers + # that interleave audio behind video (fragmented mp4) still owe most + # of it: stop only once the demuxed audio covers the cap + cap = drain_audio() + if pending_audio or samples_written >= cap: + drain_audio(final=True) + audio_done = True + break + + if output is None: + raise ValueError(f"No decodable video frames found in file '{self.__file}'") + if out_audio is not None and not audio_done: + drain_audio(final=True) + window_fill = last_video_end - last_video_pts if video_done and last_video_pts is not None else 0 + for out_packet in out_video.encode(None): + duration = video_frame_durations.pop(out_packet.pts, 0) + if out_packet.pts == last_video_pts: + duration = max(duration, window_fill) + out_packet.duration = duration + output.mux(out_packet) + if out_audio is not None: + output.mux(out_audio.encode(None)) + except BaseException: + if output is not None: + output.close() + if isinstance(path, (str, os.PathLike)) and os.path.exists(path): + os.remove(path) + raise + else: + if output is not None: + output.close() + def _get_first_video_stream(self, container: InputContainer): if len(container.streams.video): return container.streams.video[0] @@ -527,22 +801,12 @@ class VideoFromComponents(VideoInput): bit_depth: int | None = None, ): """Save the video to a file path or BytesIO buffer.""" - if format != VideoContainer.AUTO and format != VideoContainer.MP4: - raise ValueError("Only MP4 format is supported for now") - if codec != VideoCodec.AUTO and codec != VideoCodec.H264: - raise ValueError("Only H264 codec is supported for now") + open_kwargs = mp4_output_open_kwargs(path, format, codec) # None means "use the depth this video was created with" (CreateVideo's choice). if bit_depth is None: bit_depth = self.__bit_depth is_10bit = bit_depth >= 10 - extra_kwargs = {} - if isinstance(format, VideoContainer) and format != VideoContainer.AUTO: - extra_kwargs["format"] = format.value - elif isinstance(path, io.BytesIO): - # BytesIO has no file extension, so av.open can't infer the format. - # Default to mp4 since that's the only supported format anyway. - extra_kwargs["format"] = "mp4" - with av.open(path, mode='w', options={'movflags': 'use_metadata_tags'}, **extra_kwargs) as output: + with av.open(path, **open_kwargs) as output: # Add metadata before writing any streams if metadata is not None: for key, value in metadata.items(): diff --git a/tests-unit/comfy_api_test/video_types_test.py b/tests-unit/comfy_api_test/video_types_test.py index b25fcb1ca..ab4958254 100644 --- a/tests-unit/comfy_api_test/video_types_test.py +++ b/tests-unit/comfy_api_test/video_types_test.py @@ -2,11 +2,12 @@ import pytest import torch import tempfile import os +import sys import av import io from fractions import Fraction from comfy_api.input_impl.video_types import VideoFromFile, VideoFromComponents -from comfy_api.util.video_types import VideoComponents +from comfy_api.util.video_types import VideoComponents, VideoContainer, VideoCodec from comfy_api.input.basic_types import AudioInput from av.error import InvalidDataError @@ -237,3 +238,386 @@ def test_duration_consistency(video_components): manual_duration = float(components.images.shape[0] / components.frame_rate) assert duration == pytest.approx(manual_duration) + + +def create_transcode_source( + width=64, height=64, frames=30, fps=30, audio_streams=1, undecodable_audio=0, rotation=False, + container_format="mov", audio_codec="pcm_s16le", +): + """Create a temp video that save_to must transcode (mpeg4 video, so codec != h264). + + ``undecodable_audio`` trailing PCM streams get their fourcc corrupted so no decoder exists + (``codec_context is None``), like the APAC track in iPhone spatial-audio recordings. + ``rotation`` patches a 90-degree display matrix into the video track header. + """ + buffer = io.BytesIO() + with av.open(buffer, mode="w", format=container_format) as container: + video_stream = container.add_stream("mpeg4", rate=fps) + video_stream.width = width + video_stream.height = height + video_stream.pix_fmt = "yuv420p" + audio = [] + for _ in range(audio_streams + undecodable_audio): + stream = container.add_stream(audio_codec, rate=44100) + stream.sample_rate = 44100 + audio.append(stream) + + for i in range(frames): + frame = av.VideoFrame.from_ndarray( + torch.full((height, width, 3), (i * 7) % 256, dtype=torch.uint8).numpy(), + format="rgb24", + ) + container.mux(video_stream.encode(frame.reformat(format="yuv420p"))) + # write audio in 1024-sample frames, like real decoders produce, so the + # per-frame skip/cap logic in the transcode path actually runs + for stream in audio: + for offset in range(0, 44100 * frames // fps, 1024): + n = min(1024, 44100 * frames // fps - offset) + audio_frame = av.AudioFrame.from_ndarray( + torch.zeros(1, n, dtype=torch.int16).numpy(), format="s16", layout="mono" + ) + audio_frame.sample_rate = 44100 + audio_frame.pts = offset + container.mux(stream.encode(audio_frame)) + for stream in [video_stream, *audio]: + container.mux(stream.encode(None)) + + data = bytearray(buffer.getvalue()) + end = len(data) + for _ in range(undecodable_audio): + end = data.rindex(b"sowt", 0, end) + data[end:end + 4] = b"Xpac" + if rotation: + # the 3x3 display matrix sits 40 bytes into the version-0 tkhd payload; first tkhd + # inside moov = video track (search from moov so mdat bytes can't false-match) + matrix_offset = data.index(b"tkhd", data.rindex(b"moov")) + 4 + 40 + values = [0, 1 << 16, 0, -(1 << 16), 0, 0, 0, 0, 1 << 30] + data[matrix_offset:matrix_offset + 36] = b"".join(v.to_bytes(4, "big", signed=True) for v in values) + + tmp = tempfile.NamedTemporaryFile(suffix=f".{container_format}", delete=False) + tmp.write(bytes(data)) + tmp.close() + return tmp.name + + +def transcode_and_probe(video): + buffer = io.BytesIO() + video.save_to(buffer, format=VideoContainer.MP4, codec=VideoCodec.H264) + buffer.seek(0) + with av.open(buffer) as container: + video_stream = container.streams.video[0] + audio_stream = container.streams.audio[0] if container.streams.audio else None + frames = 0 + first_pts = None + for packet in container.demux(video_stream): + for frame in packet.decode(): + if first_pts is None: + first_pts = frame.pts + frames += 1 + return { + "codec": video_stream.codec_context.name, + "width": video_stream.codec_context.width, + "height": video_stream.codec_context.height, + "frames": frames, + "first_pts": first_pts, + "video_seconds": float(video_stream.duration * video_stream.time_base) if video_stream.duration else None, + "audio_seconds": float(audio_stream.duration * audio_stream.time_base) + if audio_stream and audio_stream.duration else None, + "audio_codecs": [s.codec_context.name for s in container.streams.audio], + } + + +def test_save_to_transcode_streams_without_buffering_frames(): + """Transcoding must not decode the whole video into memory first (~2 GiB for this source)""" + resource = pytest.importorskip("resource") # no getrusage on Windows + rss_scale = 1 if sys.platform == "darwin" else 1024 # ru_maxrss: bytes on macOS, KiB elsewhere + # ru_maxrss is a lifetime peak: a heavier test running earlier would shrink the measured + # delta and quietly defang this canary, so keep this source the biggest thing in the suite + file_path = create_transcode_source(width=640, height=480, frames=300) + try: + rss_before = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * rss_scale + result = transcode_and_probe(VideoFromFile(file_path)) + rss_delta = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * rss_scale - rss_before + + assert result["codec"] == "h264" + assert result["frames"] == 300 + assert rss_delta < 500 * 2**20, f"transcode buffered frames in RAM (peak grew {rss_delta / 2**20:.0f} MiB)" + finally: + os.unlink(file_path) + + +def test_save_to_transcode_honors_trim_window(): + """start_time/duration trim applies to both video and audio on the streaming path""" + file_path = create_transcode_source(frames=90) # 3s @ 30fps + try: + result = transcode_and_probe(VideoFromFile(file_path, start_time=1, duration=1)) + assert result["frames"] == pytest.approx(30, abs=2) + assert result["first_pts"] == 0 # trimmed output is rebased to start at zero + assert result["video_seconds"] == pytest.approx(1.0, abs=0.1) + assert result["audio_seconds"] == pytest.approx(1.0, abs=0.1) + finally: + os.unlink(file_path) + + +def test_save_to_transcode_keeps_audio_of_sparse_video(): + """Audio that runs ahead of a sparse video track (slideshows, timelapses) must be + kept in full — it is only clamped to the video's end, never to the video cursor.""" + buffer = io.BytesIO() + with av.open(buffer, mode="w", format="mp4") as container: + video_stream = container.add_stream("mpeg4", rate=30) + video_stream.width = video_stream.height = 64 + video_stream.pix_fmt = "yuv420p" + audio_stream = container.add_stream("aac", rate=48000, layout="stereo") + for t in (0, 30, 60): # 3 frames spread over 60 seconds + frame = av.VideoFrame.from_ndarray( + torch.full((64, 64, 3), t * 4, dtype=torch.uint8).numpy(), format="rgb24" + ).reformat(format="yuv420p") + frame.pts = t * 15360 + frame.time_base = Fraction(1, 15360) + container.mux(video_stream.encode(frame)) + container.mux(video_stream.encode(None)) + for offset in range(0, 48000 * 60, 1024): + n = min(1024, 48000 * 60 - offset) + audio_frame = av.AudioFrame.from_ndarray( + torch.zeros(2, n, dtype=torch.float32).numpy(), format="fltp", layout="stereo" + ) + audio_frame.sample_rate = 48000 + audio_frame.pts = offset + audio_frame.time_base = Fraction(1, 48000) + container.mux(audio_stream.encode(audio_frame)) + container.mux(audio_stream.encode(None)) + + buffer.seek(0) + result = transcode_and_probe(VideoFromFile(buffer)) + assert result["audio_seconds"] == pytest.approx(60.0, abs=1.0) + + +def test_save_to_transcode_vfr_audio_covers_video_span(): + """A trim window in the sparse region of a VFR file keeps audio for the true pts span + of the kept frames. Deriving the span as frames/average_rate undercuts it badly: the + average is dominated by the dense region (and can be plain wrong on MediaRecorder files).""" + buffer = io.BytesIO() + with av.open(buffer, mode="w", format="mp4") as container: + video_stream = container.add_stream("mpeg4", rate=30) + video_stream.width = video_stream.height = 64 + video_stream.pix_fmt = "yuv420p" + audio_stream = container.add_stream("aac", rate=48000, layout="stereo") + # 10 frames inside the first second, then one every 1.25 s + for i, t in enumerate([x / 10 for x in range(10)] + [1.0, 2.25, 3.5, 4.75]): + frame = av.VideoFrame.from_ndarray( + torch.full((64, 64, 3), (i * 16) % 256, dtype=torch.uint8).numpy(), format="rgb24" + ).reformat(format="yuv420p") + frame.pts = int(t * 15360) + frame.time_base = Fraction(1, 15360) + container.mux(video_stream.encode(frame)) + container.mux(video_stream.encode(None)) + for offset in range(0, 48000 * 6, 1024): + n = min(1024, 48000 * 6 - offset) + audio_frame = av.AudioFrame.from_ndarray( + torch.zeros(2, n, dtype=torch.float32).numpy(), format="fltp", layout="stereo" + ) + audio_frame.sample_rate = 48000 + audio_frame.pts = offset + audio_frame.time_base = Fraction(1, 48000) + container.mux(audio_stream.encode(audio_frame)) + container.mux(audio_stream.encode(None)) + + buffer.seek(0) + result = transcode_and_probe(VideoFromFile(buffer, start_time=1, duration=5)) + # kept frames: 1.0/2.25/3.5/4.75 s -> rebased span 3.75 s + one nominal interval + assert result["frames"] == 4 + assert result["audio_seconds"] == pytest.approx(4.0, abs=0.45) + + +def test_save_to_transcode_trims_audio_in_stream_time_base_units(): + """Matroska audio timestamps tick in 1/1000, not 1/sample_rate; trim and audio timing + must convert through the frame's time base instead of assuming sample units. AAC audio, + because it decodes straight to the encoder's format and hits the resampler passthrough + that keeps the source time base on the frames.""" + file_path = create_transcode_source(frames=90, container_format="matroska", audio_codec="aac") + try: + result = transcode_and_probe(VideoFromFile(file_path, start_time=1, duration=1)) + assert result["audio_codecs"] == ["aac"] + assert result["video_seconds"] == pytest.approx(1.0, abs=0.1) + assert result["audio_seconds"] == pytest.approx(1.0, abs=0.1) + finally: + os.unlink(file_path) + + +def test_save_to_transcode_learns_unprobed_audio_params(): + """mpegts is only probed a few seconds deep at open, so an audio stream whose first + packet comes later (live captures where audio kicks in late) still has sample_rate 0 + when the transcode starts; the parameters must be learned from the stream itself.""" + sample_rate, fps, video_seconds, audio_start = 48000, 30, 13, 12 + buffer = io.BytesIO() + with av.open(buffer, mode="w", format="mpegts") as container: + video_stream = container.add_stream("mpeg4", rate=fps) + video_stream.width = video_stream.height = 64 + video_stream.pix_fmt = "yuv420p" + audio_stream = container.add_stream("aac", rate=sample_rate, layout="mono") + for i in range(video_seconds * fps): + frame = av.VideoFrame.from_ndarray( + torch.full((64, 64, 3), (i * 7) % 256, dtype=torch.uint8).numpy(), format="rgb24" + ) + container.mux(video_stream.encode(frame.reformat(format="yuv420p"))) + for offset in range(0, (video_seconds - audio_start) * sample_rate, 1024): + n = min(1024, (video_seconds - audio_start) * sample_rate - offset) + audio_frame = av.AudioFrame.from_ndarray( + torch.zeros(1, n, dtype=torch.float32).numpy(), format="fltp", layout="mono" + ) + audio_frame.sample_rate = sample_rate + audio_frame.pts = audio_start * sample_rate + offset + container.mux(audio_stream.encode(audio_frame)) + for stream in (video_stream, audio_stream): + container.mux(stream.encode(None)) + + buffer.seek(0) + with av.open(buffer) as container: + # the scenario requires unprobed parameters; if a future FFmpeg probes deeper, + # push audio_start/video_seconds further out to restore it + assert container.streams.audio[0].codec_context.sample_rate == 0 + result = transcode_and_probe(VideoFromFile(buffer)) + assert result["frames"] == video_seconds * fps + assert result["audio_codecs"] == ["aac"] + assert result["audio_seconds"] == pytest.approx(1.0, abs=0.1) + + +def test_save_to_transcode_trimmed_fragmented_mp4_keeps_audio(): + """Fragmented mp4 (MediaRecorder, DASH/HLS-derived files) delivers audio well behind + video, so when the trim window's last video frame arrives the audio demuxed so far + does not cover the window yet; the transcode must keep demuxing audio until it does + instead of finalizing on the first audio frame it sees afterwards.""" + sample_rate, fps, seconds = 48000, 30, 6 + buffer = io.BytesIO() + with av.open(buffer, mode="w", format="mp4", options={"movflags": "frag_keyframe+empty_moov"}) as container: + video_stream = container.add_stream("h264", rate=fps) + video_stream.width = video_stream.height = 64 + video_stream.pix_fmt = "yuv420p" + audio_stream = container.add_stream("aac", rate=sample_rate, layout="mono") + next_audio_pts = 0 + for i in range(seconds * fps): + frame = av.VideoFrame.from_ndarray( + torch.full((64, 64, 3), (i * 7) % 256, dtype=torch.uint8).numpy(), format="rgb24" + ) + container.mux(video_stream.encode(frame.reformat(format="yuv420p"))) + while next_audio_pts / sample_rate <= i / fps: # feed audio alongside, like a live pipeline + audio_frame = av.AudioFrame.from_ndarray( + torch.zeros(1, 1024, dtype=torch.float32).numpy(), format="fltp", layout="mono" + ) + audio_frame.sample_rate = sample_rate + audio_frame.pts = next_audio_pts + container.mux(audio_stream.encode(audio_frame)) + next_audio_pts += 1024 + for stream in (video_stream, audio_stream): + container.mux(stream.encode(None)) + + result = transcode_and_probe(VideoFromFile(buffer, start_time=0.5, duration=1.0)) + assert result["video_seconds"] == pytest.approx(1.0, abs=0.05) + assert result["audio_seconds"] == pytest.approx(1.0, abs=0.05) + + +def test_save_to_transcode_sparse_video_keeps_true_duration(): + """average_rate is not a frame duration: a 3-frame video spanning 60 s averages + 0.05 fps, and padding the last frame with 1/average_rate used to extend the + output — and the audio kept with it — about 20 s past the source span.""" + sample_rate = 48000 + buffer = io.BytesIO() + with av.open(buffer, mode="w", format="mp4") as container: + video_stream = container.add_stream("mpeg4", rate=30) + video_stream.width = video_stream.height = 64 + video_stream.pix_fmt = "yuv420p" + audio_stream = container.add_stream("aac", rate=sample_rate, layout="mono") + for i, second in enumerate((0, 30, 60)): + frame = av.VideoFrame.from_ndarray( + torch.full((64, 64, 3), i * 80, dtype=torch.uint8).numpy(), format="rgb24" + ).reformat(format="yuv420p") + frame.pts = second * 30 + frame.time_base = Fraction(1, 30) + container.mux(video_stream.encode(frame)) + for offset in range(0, 90 * sample_rate, 1024): + n = min(1024, 90 * sample_rate - offset) + audio_frame = av.AudioFrame.from_ndarray( + torch.zeros(1, n, dtype=torch.float32).numpy(), format="fltp", layout="mono" + ) + audio_frame.sample_rate = sample_rate + audio_frame.pts = offset + container.mux(audio_stream.encode(audio_frame)) + for stream in (video_stream, audio_stream): + container.mux(stream.encode(None)) + + result = transcode_and_probe(VideoFromFile(buffer)) + assert result["frames"] == 3 + # the last frame keeps its true stts duration (1/30 s), not 1/average_rate (~20 s) + assert result["video_seconds"] == pytest.approx(60.03, abs=0.05) + assert result["audio_seconds"] == pytest.approx(60.03, abs=0.1) + + trimmed = transcode_and_probe(VideoFromFile(buffer, duration=45)) + assert trimmed["frames"] == 2 + # a kept frame whose source duration crosses the window end is clamped to it + assert trimmed["video_seconds"] == pytest.approx(45.0, abs=0.05) + assert trimmed["audio_seconds"] == pytest.approx(45.0, abs=0.1) + + +def test_save_to_transcode_irregular_vfr_keeps_span(): + """B-frames reorder packets, and mp4 sample durations follow decode order: the dts + timeline ends before the pts timeline, so an irregular-VFR source's tail holds fell + out of the container (this 20.23 s span used to come out as 15.27 s, and the 10 s + trim as 6.03 s). The transcode encodes without B-frames so every sample keeps its + true display duration.""" + durations = [1, 1, 60, 1, 1, 120, 1, 180, 1, 1, 150, 90] # 1/30 s ticks, span 20.2333 s + generator = torch.Generator().manual_seed(7) + buffer = io.BytesIO() + with av.open(buffer, mode="w", format="mp4") as container: + video_stream = container.add_stream("mpeg4", rate=30) + video_stream.width = video_stream.height = 64 + video_stream.pix_fmt = "yuv420p" + pts = 0 + for duration in durations: + # textured frames, so an encoder with default settings has B-frames to gain from + frame = av.VideoFrame.from_ndarray( + torch.randint(0, 255, (64, 64, 3), generator=generator, dtype=torch.uint8).numpy(), + format="rgb24", + ).reformat(format="yuv420p") + frame.pts = pts + frame.time_base = Fraction(1, 30) + pts += duration + for packet in video_stream.encode(frame): + packet.duration = duration # exact stts in the source + container.mux(packet) + container.mux(video_stream.encode(None)) + + result = transcode_and_probe(VideoFromFile(buffer)) + assert result["frames"] == len(durations) + assert result["video_seconds"] == pytest.approx(sum(durations) / 30, abs=0.05) + + trimmed = transcode_and_probe(VideoFromFile(buffer, duration=10)) + assert trimmed["frames"] == 8 # frames at 12.167 s+ fall outside the window + assert trimmed["video_seconds"] == pytest.approx(10.0, abs=0.05) + + +def test_save_to_transcode_bakes_rotation(): + """A 90-degree display-matrix rotation swaps the output dimensions (portrait video)""" + file_path = create_transcode_source(width=64, height=32, rotation=True) + try: + result = transcode_and_probe(VideoFromFile(file_path)) + assert (result["width"], result["height"]) == (32, 64) + assert result["frames"] == 30 + finally: + os.unlink(file_path) + + +def test_save_to_transcode_skips_undecodable_audio(): + """Streaming transcode keeps the decodable audio track and drops undecodable ones; + with no decodable audio at all the output is video-only instead of crashing.""" + mixed = all_bad = None + try: + mixed = create_transcode_source(audio_streams=1, undecodable_audio=1) + all_bad = create_transcode_source(audio_streams=0, undecodable_audio=2) + result = transcode_and_probe(VideoFromFile(mixed)) + assert result["audio_codecs"] == ["aac"] + assert result["audio_seconds"] == pytest.approx(1.0, abs=0.1) + assert transcode_and_probe(VideoFromFile(all_bad))["audio_codecs"] == [] + finally: + for path in (mixed, all_bad): + if path: + os.unlink(path) From be27c9c3df4551da2f3df9b57b54d3b33b5cf36c Mon Sep 17 00:00:00 2001 From: bigcat88 Date: Wed, 8 Jul 2026 15:01:57 +0300 Subject: [PATCH 2/3] fix: trim transcode with missing leading PTS Signed-off-by: bigcat88 --- comfy_api/latest/_input_impl/video_types.py | 3 +- tests-unit/comfy_api_test/video_types_test.py | 55 +++++++++++++++++++ 2 files changed, 57 insertions(+), 1 deletion(-) diff --git a/comfy_api/latest/_input_impl/video_types.py b/comfy_api/latest/_input_impl/video_types.py index b9511e593..3ab8fa2d0 100644 --- a/comfy_api/latest/_input_impl/video_types.py +++ b/comfy_api/latest/_input_impl/video_types.py @@ -626,7 +626,8 @@ class VideoFromFile(VideoInput): video_done = True if last_video_pts is not None: # the source continues past the window: hold the last kept frame to the window end - last_video_end = max(last_video_end, end_pts - video_pts_offset) + end_offset = video_pts_offset if video_pts_offset is not None else start_pts + last_video_end = max(last_video_end, end_pts - end_offset) break # the source's true display duration of this frame; average_rate is not a # frame duration (sparse/VFR sources), so it is only the fallback diff --git a/tests-unit/comfy_api_test/video_types_test.py b/tests-unit/comfy_api_test/video_types_test.py index ab4958254..9d78b28b3 100644 --- a/tests-unit/comfy_api_test/video_types_test.py +++ b/tests-unit/comfy_api_test/video_types_test.py @@ -595,6 +595,61 @@ def test_save_to_transcode_irregular_vfr_keeps_span(): assert trimmed["video_seconds"] == pytest.approx(10.0, abs=0.05) +def test_save_to_transcode_trim_survives_missing_leading_pts(): + """A trim should survive pts-less kept frames followed by a real-pts frame past the window.""" + nulled_frames = 0 + + class _PacketProxy: + def __init__(self, packet): + self._packet = packet + + def __getattr__(self, name): + return getattr(self._packet, name) + + @property + def stream(self): + return self._packet.stream + + def decode(self): + nonlocal nulled_frames + frames = self._packet.decode() + for frame in frames: + if nulled_frames < 2: + frame.pts = None + nulled_frames += 1 + return frames + + class _ContainerProxy: + def __init__(self, real): + self._real = real + + def __getattr__(self, name): + return getattr(self._real, name) + + def demux(self, *streams): + for packet in self._real.demux(*streams): + yield _PacketProxy(packet) + + file_path = create_transcode_source(frames=10, audio_streams=0) + try: + buffer = io.BytesIO() + with av.open(file_path) as container: + # 0.05 s window: both pts-less frames are kept (synthesized pts 0 and 512), + # and the first real-pts frame (1024 ticks) already lies past end_pts (768) + VideoFromFile(file_path, duration=0.05)._save_transcoded( + _ContainerProxy(container), buffer, VideoContainer.MP4, VideoCodec.H264, None, 8 + ) + assert nulled_frames == 2 + buffer.seek(0) + with av.open(buffer) as container: + video_stream = container.streams.video[0] + frames = [f for p in container.demux(video_stream) for f in p.decode()] + assert len(frames) == 2 + assert float(video_stream.duration * video_stream.time_base) == pytest.approx(2 / 30, abs=0.01) + finally: + os.unlink(file_path) + + def test_save_to_transcode_bakes_rotation(): """A 90-degree display-matrix rotation swaps the output dimensions (portrait video)""" file_path = create_transcode_source(width=64, height=32, rotation=True) From 64f73d40799808431a2e2c900025ece6a60a6f4f Mon Sep 17 00:00:00 2001 From: bigcat88 Date: Wed, 8 Jul 2026 15:35:32 +0300 Subject: [PATCH 3/3] fix: continue probing audio parameters after a packet decode error instead of dropping stream immediately Signed-off-by: bigcat88 --- comfy_api/latest/_input_impl/video_types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/comfy_api/latest/_input_impl/video_types.py b/comfy_api/latest/_input_impl/video_types.py index 3ab8fa2d0..9f498a6bb 100644 --- a/comfy_api/latest/_input_impl/video_types.py +++ b/comfy_api/latest/_input_impl/video_types.py @@ -80,7 +80,7 @@ def probe_audio_params(container: InputContainer, audio_stream, max_packets: int try: frames = packet.decode() except av.error.FFmpegError: - return 0, 0 + frames = () if frames: return frames[0].sample_rate, frames[0].layout.nb_channels if i >= max_packets: