mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-09-05 23:47:12 +08:00
Fix video trimming corruption with decode/re-encode approach (#225)
This commit is contained in:
parent
be206ef4de
commit
304606a645
@ -221,39 +221,69 @@ def trim_video(video: VideoInput, duration_sec: float) -> VideoInput:
|
|||||||
input_container = av.open(input_source, mode='r')
|
input_container = av.open(input_source, mode='r')
|
||||||
output_container = av.open(output_buffer, mode='w', format='mp4')
|
output_container = av.open(output_buffer, mode='w', format='mp4')
|
||||||
|
|
||||||
# Set up stream mapping
|
# Set up output streams for re-encoding
|
||||||
stream_map = {}
|
video_stream = None
|
||||||
|
audio_stream = None
|
||||||
|
|
||||||
for stream in input_container.streams:
|
for stream in input_container.streams:
|
||||||
if stream.type in ('video', 'audio'):
|
logging.info(f"Found stream: type={stream.type}, class={type(stream)}")
|
||||||
out_stream = output_container.add_stream_from_template(template=stream)
|
if isinstance(stream, av.VideoStream):
|
||||||
stream_map[stream] = out_stream
|
# Create output video stream with same parameters
|
||||||
|
video_stream = output_container.add_stream('h264', rate=stream.average_rate)
|
||||||
|
video_stream.width = stream.width
|
||||||
|
video_stream.height = stream.height
|
||||||
|
video_stream.pix_fmt = 'yuv420p'
|
||||||
|
logging.info(f"Added video stream: {stream.width}x{stream.height} @ {stream.average_rate}fps")
|
||||||
|
elif isinstance(stream, av.AudioStream):
|
||||||
|
# Create output audio stream with same parameters
|
||||||
|
audio_stream = output_container.add_stream('aac', rate=stream.sample_rate)
|
||||||
|
audio_stream.sample_rate = stream.sample_rate
|
||||||
|
audio_stream.layout = stream.layout
|
||||||
|
logging.info(f"Added audio stream: {stream.sample_rate}Hz, {stream.channels} channels")
|
||||||
|
|
||||||
# Since we're always starting from 0, no need to seek
|
frame_count = 0
|
||||||
# Just process packets until we reach end_sec
|
audio_frame_count = 0
|
||||||
|
|
||||||
for packet in input_container.demux():
|
# Decode and re-encode video frames
|
||||||
if packet.stream not in stream_map:
|
if video_stream:
|
||||||
continue
|
for frame in input_container.decode(video=0):
|
||||||
|
if frame.time >= duration_sec:
|
||||||
# Get packet timestamp (prefer PTS, fallback to DTS)
|
|
||||||
pts = packet.pts if packet.pts is not None else packet.dts
|
|
||||||
if pts is None:
|
|
||||||
continue # Skip packets without timestamps
|
|
||||||
|
|
||||||
time_in_seconds = float(pts * packet.time_base)
|
|
||||||
|
|
||||||
# Stop when we reach the target duration
|
|
||||||
if time_in_seconds >= duration_sec:
|
|
||||||
break
|
break
|
||||||
|
|
||||||
# Remap packet to output stream (timestamps already start at 0)
|
# Re-encode frame
|
||||||
packet.stream = stream_map[packet.stream]
|
for packet in video_stream.encode(frame):
|
||||||
output_container.mux(packet)
|
output_container.mux(packet)
|
||||||
|
frame_count += 1
|
||||||
|
|
||||||
|
# Flush encoder
|
||||||
|
for packet in video_stream.encode():
|
||||||
|
output_container.mux(packet)
|
||||||
|
|
||||||
|
logging.info(f"Encoded {frame_count} video frames")
|
||||||
|
|
||||||
|
# Decode and re-encode audio frames
|
||||||
|
if audio_stream:
|
||||||
|
input_container.seek(0) # Reset to beginning for audio
|
||||||
|
for frame in input_container.decode(audio=0):
|
||||||
|
if frame.time >= duration_sec:
|
||||||
|
break
|
||||||
|
|
||||||
|
# Re-encode frame
|
||||||
|
for packet in audio_stream.encode(frame):
|
||||||
|
output_container.mux(packet)
|
||||||
|
audio_frame_count += 1
|
||||||
|
|
||||||
|
# Flush encoder
|
||||||
|
for packet in audio_stream.encode():
|
||||||
|
output_container.mux(packet)
|
||||||
|
|
||||||
|
logging.info(f"Encoded {audio_frame_count} audio frames")
|
||||||
|
|
||||||
# Close containers
|
# Close containers
|
||||||
output_container.close()
|
output_container.close()
|
||||||
input_container.close()
|
input_container.close()
|
||||||
|
|
||||||
|
|
||||||
# Return as VideoFromFile using the buffer
|
# Return as VideoFromFile using the buffer
|
||||||
output_buffer.seek(0)
|
output_buffer.seek(0)
|
||||||
return VideoFromFile(output_buffer)
|
return VideoFromFile(output_buffer)
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user