mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-08 21:36:14 +08:00
Compare commits
1 Commits
master
...
fix/core/v
| Author | SHA1 | Date | |
|---|---|---|---|
| 1d595d4a70 |
@ -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():
|
||||
|
||||
@ -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)
|
||||
|
||||
Reference in New Issue
Block a user