mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-08-24 01:56:59 +08:00
Keep subtitles when the transcode path runs
_save_transcoded had no subtitle output path, so every stream was dropped whenever save_to re-encoded: both on the explicit format/codec/trim route and on the fallback out of _save_remuxed. Compatible subtitles survived a remux and vanished the moment anything forced a transcode. Subtitle streams are now carried through untouched, rebased onto the same trimmed timeline as the video and clipped to the trim window. There is no subtitle encoder binding in PyAV, so a stream the output container cannot store as-is is still dropped -- but named in a warning, matching what _save_remuxed already does rather than disappearing silently. The remux fallback warning no longer claims subtitles are dropped, because the transcode it hands off to now keeps the ones it can.
This commit is contained in:
@@ -538,7 +538,8 @@ class VideoFromFile(VideoInput):
|
||||
continue
|
||||
logging.warning(
|
||||
"The %s container cannot store %s, so the whole file is being re-encoded to H.264/AAC. "
|
||||
"Any additional audio or subtitle streams will be dropped.",
|
||||
"Any additional audio streams will be dropped; subtitles the output container "
|
||||
"can store are kept.",
|
||||
format_name, codec_name,
|
||||
)
|
||||
return False
|
||||
@@ -602,7 +603,16 @@ class VideoFromFile(VideoInput):
|
||||
if duration:
|
||||
duration_cap = math.ceil(duration * sample_rate)
|
||||
|
||||
# Subtitles are remuxed untouched: there is no subtitle encoder binding, so a stream the
|
||||
# output container cannot store as-is is dropped with a warning naming it, exactly like
|
||||
# the remux path does. Streams FFmpeg has no decoder for cannot template a new stream.
|
||||
subtitle_streams = [s for s in container.streams.subtitles if s.codec_context is not None]
|
||||
streams = [video_stream] if audio_stream is None else [video_stream, audio_stream]
|
||||
streams += subtitle_streams
|
||||
subtitle_map = {}
|
||||
# Subtitle packets that arrive before the first kept video frame: the output is not open
|
||||
# yet and the pts rebase offset is not known, so they wait here rather than being lost.
|
||||
pending_subtitles = []
|
||||
pts_step = max(1, int(round((1 / rate) / video_stream.time_base)))
|
||||
video_done = False
|
||||
audio_done = audio_stream is None
|
||||
@@ -660,6 +670,32 @@ class VideoFromFile(VideoInput):
|
||||
audio_done = True
|
||||
return cap
|
||||
|
||||
def mux_subtitle(packet):
|
||||
"""Remux one subtitle packet, rebased onto the trimmed timeline the video was rebased to."""
|
||||
out_stream = subtitle_map.get(packet.stream)
|
||||
if out_stream is None or packet.dts is None or packet.pts is None or packet.time_base is None:
|
||||
return
|
||||
start = float(packet.pts * packet.time_base)
|
||||
if start < start_time or (duration and start >= start_time + duration):
|
||||
return
|
||||
# the video's own rebase offset, so subtitles stay in sync with it rather than
|
||||
# with the requested start (a seek lands on the preceding keyframe)
|
||||
offset_ticks = video_pts_offset if video_pts_offset is not None else start_pts
|
||||
shift = int(round(float(offset_ticks * video_stream.time_base) / packet.time_base))
|
||||
packet.pts -= shift
|
||||
packet.dts -= shift
|
||||
if packet.pts < 0:
|
||||
return
|
||||
packet.stream = out_stream
|
||||
output.mux(packet)
|
||||
|
||||
def flush_subtitles():
|
||||
# only once the first video frame fixed the rebase offset
|
||||
if output is None or not pending_subtitles or last_video_pts is None:
|
||||
return
|
||||
while pending_subtitles:
|
||||
mux_subtitle(pending_subtitles.pop(0))
|
||||
|
||||
try:
|
||||
for packet in container.demux(*streams):
|
||||
if video_done and audio_done:
|
||||
@@ -711,6 +747,19 @@ class VideoFromFile(VideoInput):
|
||||
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)
|
||||
for subtitle_stream in subtitle_streams:
|
||||
try:
|
||||
subtitle_map[subtitle_stream] = output.add_stream_from_template(
|
||||
template=subtitle_stream, opaque=True
|
||||
)
|
||||
except ValueError:
|
||||
logging.warning(
|
||||
"Dropping %s subtitle stream %d: the %s container cannot store it, "
|
||||
"and subtitles cannot be re-encoded.",
|
||||
subtitle_stream.codec_context.name,
|
||||
subtitle_stream.index,
|
||||
output.format.name,
|
||||
)
|
||||
if (frame.width, frame.height) != source_size:
|
||||
# encoding would silently rescale the new geometry into the old one
|
||||
raise ValueError(
|
||||
@@ -772,6 +821,7 @@ class VideoFromFile(VideoInput):
|
||||
for out_packet in out_video.encode(frame):
|
||||
out_packet.duration = video_frame_durations.pop(out_packet.pts, 0)
|
||||
output.mux(out_packet)
|
||||
flush_subtitles()
|
||||
drain_audio()
|
||||
|
||||
elif packet.stream == audio_stream and not audio_done:
|
||||
@@ -807,6 +857,12 @@ class VideoFromFile(VideoInput):
|
||||
audio_done = True
|
||||
break
|
||||
|
||||
elif packet.stream in subtitle_map:
|
||||
mux_subtitle(packet)
|
||||
elif packet.stream in subtitle_streams and output is None:
|
||||
pending_subtitles.append(packet)
|
||||
|
||||
flush_subtitles()
|
||||
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:
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import logging
|
||||
import pytest
|
||||
import torch
|
||||
import tempfile
|
||||
@@ -744,6 +745,7 @@ def test_save_to_transcode_clamps_final_pts_to_declared_stream_duration():
|
||||
def __init__(self, video_stream):
|
||||
self.video = [video_stream]
|
||||
self.audio = []
|
||||
self.subtitles = []
|
||||
|
||||
class _PacketProxy:
|
||||
def __init__(self, packet, stream):
|
||||
@@ -888,6 +890,139 @@ def test_save_to_transcode_bakes_rotation():
|
||||
os.unlink(file_path)
|
||||
|
||||
|
||||
def mov_text_payload(text: str) -> bytes:
|
||||
"""A mov_text sample is a 16-bit big-endian length followed by the UTF-8 text."""
|
||||
encoded = text.encode("utf-8")
|
||||
return len(encoded).to_bytes(2, "big") + encoded
|
||||
|
||||
|
||||
def create_subtitled_source(subtitle_codec="mov_text", container_format="mp4", frames=90, fps=30):
|
||||
"""mpeg4 video (so save_to must transcode) alongside a subtitle track carrying three cues."""
|
||||
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 = video_stream.height = 64
|
||||
video_stream.pix_fmt = "yuv420p"
|
||||
subtitle_stream = container.add_mux_stream(subtitle_codec)
|
||||
subtitle_stream.time_base = Fraction(1, 1000)
|
||||
for i in range(frames):
|
||||
frame = av.VideoFrame.from_ndarray(
|
||||
torch.full((64, 64, 3), (i * 7) % 256, dtype=torch.uint8).numpy(), format="rgb24"
|
||||
).reformat(format="yuv420p")
|
||||
container.mux(video_stream.encode(frame))
|
||||
container.mux(video_stream.encode(None))
|
||||
for start_ms, text in ((0, "one"), (1000, "two"), (2000, "three")):
|
||||
packet = av.Packet(mov_text_payload(text))
|
||||
packet.stream = subtitle_stream
|
||||
packet.pts = packet.dts = start_ms
|
||||
packet.duration = 900
|
||||
packet.time_base = Fraction(1, 1000)
|
||||
container.mux(packet)
|
||||
buffer.seek(0)
|
||||
return buffer
|
||||
|
||||
|
||||
def subtitle_cues(buffer):
|
||||
"""(seconds, text) for every non-empty cue; the mp4 muxer pads gaps with 2-byte empties."""
|
||||
buffer.seek(0)
|
||||
with av.open(buffer) as container:
|
||||
if not container.streams.subtitles:
|
||||
return None
|
||||
return [
|
||||
(float(packet.pts * packet.time_base), bytes(packet)[2:].decode("utf-8"))
|
||||
for packet in container.demux(container.streams.subtitles[0])
|
||||
if packet.dts is not None and packet.size > 2
|
||||
]
|
||||
|
||||
|
||||
def test_save_to_transcode_keeps_subtitles_the_container_can_store():
|
||||
"""Transcoding video must not silently drop a subtitle track the output can hold."""
|
||||
output = io.BytesIO()
|
||||
VideoFromFile(create_subtitled_source()).save_to(
|
||||
output, format=VideoContainer.MP4, codec=VideoCodec.H264
|
||||
)
|
||||
|
||||
output.seek(0)
|
||||
with av.open(output) as container:
|
||||
assert container.streams.video[0].codec_context.name == "h264"
|
||||
assert [s.codec_context.name for s in container.streams.subtitles] == ["mov_text"]
|
||||
assert subtitle_cues(output) == [(0.0, "one"), (1.0, "two"), (2.0, "three")]
|
||||
|
||||
|
||||
def test_save_to_transcode_trims_subtitles_with_the_video():
|
||||
"""Kept cues rebase onto the trimmed timeline; cues outside the window are dropped."""
|
||||
output = io.BytesIO()
|
||||
VideoFromFile(create_subtitled_source(), start_time=1, duration=1).save_to(
|
||||
output, format=VideoContainer.MP4, codec=VideoCodec.H264
|
||||
)
|
||||
|
||||
assert subtitle_cues(output) == [(0.0, "two")]
|
||||
|
||||
|
||||
def test_save_to_transcode_drops_unstorable_subtitles_with_a_warning(caplog):
|
||||
"""There is no subtitle encoder binding, so subrip cannot become mov_text: drop it, but
|
||||
name the stream instead of letting the track vanish silently."""
|
||||
output = io.BytesIO()
|
||||
with caplog.at_level(logging.WARNING):
|
||||
VideoFromFile(create_subtitled_source(subtitle_codec="subrip", container_format="matroska")).save_to(
|
||||
output, format=VideoContainer.MP4, codec=VideoCodec.H264
|
||||
)
|
||||
|
||||
assert subtitle_cues(output) is None
|
||||
assert any(
|
||||
"subtitle stream" in record.message and "cannot store it" in record.message
|
||||
for record in caplog.records
|
||||
)
|
||||
output.seek(0)
|
||||
with av.open(output) as container:
|
||||
assert container.streams.video[0].codec_context.name == "h264"
|
||||
assert len([f for p in container.demux(container.streams.video[0]) for f in p.decode()]) == 90
|
||||
|
||||
|
||||
def test_save_to_remux_fallback_keeps_subtitles():
|
||||
"""The audio-triggered fallback into the transcode path must keep subtitles too."""
|
||||
buffer = io.BytesIO()
|
||||
with av.open(buffer, mode="w", format="mov") 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("pcm_u8", rate=44100)
|
||||
audio_stream.sample_rate = 44100
|
||||
subtitle_stream = container.add_mux_stream("mov_text")
|
||||
subtitle_stream.time_base = Fraction(1, 1000)
|
||||
for i in range(30):
|
||||
frame = av.VideoFrame.from_ndarray(
|
||||
torch.full((64, 64, 3), (i * 7) % 256, dtype=torch.uint8).numpy(), format="rgb24"
|
||||
).reformat(format="yuv420p")
|
||||
container.mux(video_stream.encode(frame))
|
||||
for offset in range(0, 44100, 1024):
|
||||
audio_frame = av.AudioFrame.from_ndarray(
|
||||
torch.zeros(1, min(1024, 44100 - offset), dtype=torch.int16).numpy(),
|
||||
format="s16", layout="mono",
|
||||
)
|
||||
audio_frame.sample_rate = 44100
|
||||
audio_frame.pts = offset
|
||||
container.mux(audio_stream.encode(audio_frame))
|
||||
for stream in (video_stream, audio_stream):
|
||||
container.mux(stream.encode(None))
|
||||
packet = av.Packet(mov_text_payload("one"))
|
||||
packet.stream = subtitle_stream
|
||||
packet.pts = packet.dts = 0
|
||||
packet.duration = 900
|
||||
packet.time_base = Fraction(1, 1000)
|
||||
container.mux(packet)
|
||||
|
||||
buffer.seek(0)
|
||||
output = io.BytesIO()
|
||||
VideoFromFile(buffer).save_to(output, format=VideoContainer.MP4)
|
||||
|
||||
output.seek(0)
|
||||
with av.open(output) as container:
|
||||
assert container.streams.video[0].codec_context.name == "h264"
|
||||
assert container.streams.audio[0].codec_context.name == "aac"
|
||||
assert subtitle_cues(output) == [(0.0, "one")]
|
||||
|
||||
|
||||
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."""
|
||||
|
||||
Reference in New Issue
Block a user