From 313400f778f6ee222cffda30e17273736821437f Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Mon, 12 Oct 2026 21:29:14 -0400 Subject: [PATCH 1/7] fix(video): stop re-encoding earlier segments when extending videos video_concat and extract_video_range stream-copy H.264 frames and AAC packets the app already encoded, re-encoding only the frames and audio around each cut, with lossless splices verified and a full re-encode as fallback. Video encodes now use CRF 18 with a 2 s IDR interval so later joins have cut points. --- .../05.Video Generation/introduction.mdx | 2 +- invokeai/app/invocations/video_concat.py | 422 +++++++++++++---- .../invocations/video_frame_extract_range.py | 199 +++++++- invokeai/app/util/aac_splice.py | 359 ++++++++++++++ invokeai/app/util/h264_splice.py | 444 ++++++++++++++++++ invokeai/app/util/mp4_metadata.py | 24 +- invokeai/app/util/video_audio.py | 6 +- invokeai/app/util/video_encoding.py | 59 ++- invokeai/app/util/video_thumbnails.py | 4 +- tests/app/invocations/test_video_concat.py | 4 +- .../invocations/test_video_concat_splice.py | 360 ++++++++++++++ .../test_video_frame_extract_range.py | 206 +++++++- tests/app/util/test_h264_splice.py | 103 ++++ 13 files changed, 2063 insertions(+), 129 deletions(-) create mode 100644 invokeai/app/util/aac_splice.py create mode 100644 invokeai/app/util/h264_splice.py create mode 100644 tests/app/invocations/test_video_concat_splice.py create mode 100644 tests/app/util/test_h264_splice.py diff --git a/docs/src/content/docs/Users Guide/05.Video Generation/introduction.mdx b/docs/src/content/docs/Users Guide/05.Video Generation/introduction.mdx index f375d78efe9..3cf8efa3d2b 100644 --- a/docs/src/content/docs/Users Guide/05.Video Generation/introduction.mdx +++ b/docs/src/content/docs/Users Guide/05.Video Generation/introduction.mdx @@ -203,7 +203,7 @@ Extension is available on all three families. LTX-2's continues the soundtrack a ### Quality drift across iterations -Each iteration's source is itself a generation output, so artifacts compound — by the 4th or 5th extension you may see softening or color drift. Mitigations: +Each iteration's source is itself a generation output, so artifacts compound — by the 4th or 5th extension you may see softening or color drift. That drift is in the *new* footage: the segments already in the initial video are carried into the result without being re-encoded, picture and sound, apart from up to two seconds of picture at each cut — before the join, and at a trim's start. A trimmed initial video's soundtrack is re-encoded once. Extending a clip whose frame rate differs from the model's (MiniMax H3 always runs at 24 fps) re-encodes it in full, once, as does extending an uploaded video or one made by an earlier version of Invoke. Mitigations: 1. **Trim the bridge point a few frames back from the end** — boundary frames are the most artifact-prone. 2. **Refresh a bridge frame with a low-strength img2img pass** (SDXL or FLUX at ~0.2 strength) before using it as a first frame for the next segment. diff --git a/invokeai/app/invocations/video_concat.py b/invokeai/app/invocations/video_concat.py index 1f9537f4c4a..8ddeff8ba0b 100644 --- a/invokeai/app/invocations/video_concat.py +++ b/invokeai/app/invocations/video_concat.py @@ -11,6 +11,16 @@ O(transition_frames) even when the inputs are long uploads (the upload cap admits files whose decoded frames would run to tens of gigabytes). +The first clip is not re-encoded when it can be avoided: if it was written by this +app's own encoder at the output canvas and frame rate, everything before its last +usable keyframe is stream-copied, and only the remainder — its tail, the transitions +and the later clips — is encoded and joined on (see ``h264_splice``). Repeatedly +extending a video therefore stops degrading its earlier segments, which otherwise go +through one more lossy generation per extension. Any other first clip is re-encoded +in full, as before. The soundtrack is treated the same way: the first clip's own AAC +packets are kept up to its outgoing transition and only the rest is encoded (see +``aac_splice``). + The AUDIO path is a deliberate exception to that bound: soundtracks buffer in full (float32 stereo PCM of the emitted timeline, roughly 23 MB per minute at 48 kHz). Audio is ~3 orders of magnitude smaller per second than decoded @@ -23,6 +33,7 @@ import math import tempfile from collections import deque +from itertools import islice from pathlib import Path from typing import Callable, Iterable, Iterator, Literal, Optional @@ -39,8 +50,28 @@ from invokeai.app.invocations.primitives import VideoOutput from invokeai.app.services.session_processor.session_processor_common import CanceledException from invokeai.app.services.shared.invocation_context import InvocationContext -from invokeai.app.util.video_audio import extract_audio_pcm, mux_audio_into_video, resample_linear -from invokeai.app.util.video_encoding import make_mp4_writer, write_stereo_wav +from invokeai.app.util.aac_splice import ( + AAC_FRAME_SAMPLES, + PREROLL_FRAMES, + encoder_aac_layout, + mux_spliced_aac, + read_aac_layout, +) +from invokeai.app.util.h264_splice import ( + MIN_ENCODED_FRAMES, + H264Layout, + SpliceError, + concat_stream_copy, + copy_frames, + read_h264_layout, +) +from invokeai.app.util.video_audio import ( + AudioExtractionError, + extract_audio_pcm, + mux_audio_into_video, + resample_linear, +) +from invokeai.app.util.video_encoding import make_mp4_writer, write_stereo_wav, writer_decoder_config from invokeai.app.util.video_thumbnails import iter_video_frames, probe_video TransitionMode = Literal["cut", "crossfade", "fade_through_black"] @@ -195,84 +226,302 @@ def invoke(self, context: InvocationContext) -> VideoOutput: context.util.signal_progress(f"Joining {len(self.videos)} clip(s) ({self.transition}) @ {output_fps:.2f} fps") - tmp = tempfile.NamedTemporaryFile(prefix="invokeai_video_concat_", suffix=".mp4", delete=False) - tmp.close() - tmp_path = Path(tmp.name) - wav_path: Path | None = None - muxed_path: Path | None = None - try: - # Frames stream from the decoders straight into the encoder; only the - # transition windows are buffered. See _iter_joined_frames. - writer = make_mp4_writer(tmp_path, output_fps) - num_frames = 0 - frame_counts: list[int] = [] - try: - # Every clip goes through the canvas guard, including the first: frames already - # on it pass straight through, so nothing can reach the writer at a size the - # file is not being encoded at. - clip_iters = [ - _iter_frames_on_canvas(iter_video_frames(p, is_canceled=context.util.is_canceled), (width, height)) - for p in paths - ] - for frame in self._iter_joined_frames( - clip_iters, is_canceled=context.util.is_canceled, frame_counts=frame_counts - ): - writer.append_data(frame) - num_frames += 1 - finally: - writer.close() - - if num_frames == 0: - raise ValueError("Concatenation produced zero output frames.") - - # Rebuild the soundtrack on the emitted timeline. Inputs without an audio - # track (e.g. Wan clips) contribute silence; if none carries audio, the - # output has no audio stream, as before. - source_path = tmp_path - audio = self._build_audio_track( - context=context, - paths=paths, - frame_counts=frame_counts, - native_rates=[probe[3] for probe in probes], - native_durations=[probe[2] for probe in probes], - output_fps=output_fps, - num_output_frames=num_frames, - ) - if audio is not None: - pcm, rate = audio - context.util.signal_progress("Muxing audio") - wav_tmp = tempfile.NamedTemporaryFile(prefix="invokeai_video_concat_", suffix=".wav", delete=False) - wav_tmp.close() - wav_path = Path(wav_tmp.name) + with tempfile.TemporaryDirectory(prefix="invokeai_video_concat_") as work: + work_dir = Path(work) + tmp_path = work_dir / "joined.mp4" + joined = None + splice = self._plan_splice(paths[0], probes[0][3], (width, height), output_fps) + if splice is not None: + keyframe, layout = splice + rest_path = work_dir / "rest.mp4" + rest = self._encode_joined(context, paths, (width, height), output_fps, rest_path, skip=keyframe) + if self._splice(context, paths[0], keyframe, layout, rest_path, rest[0][0], work_dir, tmp_path): + joined = rest + if joined is None: + joined = self._encode_joined(context, paths, (width, height), output_fps, tmp_path, skip=0) + frame_counts, num_frames = joined + return self._save_joined(context, paths, probes, tmp_path, work_dir, frame_counts, num_frames, output_fps) + + def _save_joined( + self, + context: InvocationContext, + paths: list[Path], + probes: list[tuple[int, int, Optional[float], Optional[float]]], + tmp_path: Path, + work_dir: Path, + frame_counts: list[int], + num_frames: int, + output_fps: float, + ) -> VideoOutput: + width, height = probes[0][0], probes[0][1] + if num_frames == 0: + raise ValueError("Concatenation produced zero output frames.") + + # Rebuild the soundtrack on the emitted timeline. Inputs without an audio + # track (e.g. Wan clips) contribute silence; if none carries audio, the + # output has no audio stream, as before. + source_path = tmp_path + audio = self._build_audio_track( + context=context, + paths=paths, + frame_counts=frame_counts, + native_rates=[probe[3] for probe in probes], + native_durations=[probe[2] for probe in probes], + output_fps=output_fps, + num_output_frames=num_frames, + ) + if audio is not None: + pcm, rate = audio + context.util.signal_progress("Muxing audio") + muxed_path = work_dir / "muxed.mp4" + if not self._mux_spliced_audio( + context, paths[0], probes[0][3], frame_counts[0], pcm, rate, output_fps, tmp_path, muxed_path + ): + wav_path = work_dir / "audio.wav" write_stereo_wav(wav_path, pcm, rate) - mux_tmp = tempfile.NamedTemporaryFile(prefix="invokeai_video_concat_", suffix=".mp4", delete=False) - mux_tmp.close() - muxed_path = Path(mux_tmp.name) mux_audio_into_video(tmp_path, wav_path, muxed_path) - source_path = muxed_path + source_path = muxed_path + + duration = num_frames / output_fps + context.logger.info( + f"Encoded concatenated MP4: {num_frames} frames @ {output_fps:.2f} fps " + f"({duration:.2f}s) at {width}x{height}" + (", with AAC audio" if audio is not None else ", silent") + ) + video_dto = context.videos.save( + source_path=source_path, + width=width, + height=height, + duration=duration, + fps=output_fps, + ) + context.logger.info(f"Saved concatenated video: {video_dto.video_name}") + return VideoOutput.build(video_dto) + + def _transition_windows(self) -> tuple[int, int]: + """``(tail_need, head_need)``: frames each boundary takes from the outgoing clip's + tail and the incoming clip's head. The class docstring describes the layout.""" + if self.transition == "crossfade" and self.transition_frames > 0: + return self.transition_frames, self.transition_frames + if self.transition == "fade_through_black" and self.transition_frames > 0: + tail_need = self.transition_frames // 2 + return tail_need, self.transition_frames - tail_need + return 0, 0 + + def _plan_splice( + self, first_path: Path, first_fps: Optional[float], canvas: tuple[int, int], output_fps: float + ) -> Optional[tuple[int, H264Layout]]: + """Pick the keyframe of the first clip up to which it can be stream-copied, or None. + + The keyframe must leave the clip's transition tail, and at least ``MIN_ENCODED_FRAMES`` + so the re-encoded remainder keeps the copied head's timeline, after it. The clip must + have a constant frame duration and already play at the + output rate — a stream copy cannot retime — and carry the decoder configuration this + node's encoder writes, or the re-encoded remainder could not share its track. + """ + if first_fps is None or not math.isclose(first_fps, output_fps, rel_tol=1e-6): + return None + layout = read_h264_layout(first_path) + if ( + layout is None + or layout.frame_duration is None + or layout.decoder_config != writer_decoder_config(*canvas, output_fps) + ): + return None + tail_need, _ = self._transition_windows() + last_usable = layout.sample_count - max(tail_need, MIN_ENCODED_FRAMES) + keyframes = [k for k in layout.sync_samples if 0 < k <= last_usable] + return (keyframes[-1], layout) if keyframes else None - duration = num_frames / output_fps - context.logger.info( - f"Encoded concatenated MP4: {num_frames} frames @ {output_fps:.2f} fps " - f"({duration:.2f}s) at {width}x{height}" + (", with AAC audio" if audio is not None else ", silent") - ) - video_dto = context.videos.save( - source_path=source_path, - width=width, - height=height, - duration=duration, - fps=output_fps, - ) - context.logger.info(f"Saved concatenated video: {video_dto.video_name}") - return VideoOutput.build(video_dto) + def _encode_joined( + self, + context: InvocationContext, + paths: list[Path], + canvas: tuple[int, int], + output_fps: float, + out_path: Path, + skip: int, + ) -> tuple[list[int], int]: + """Encode the joined timeline, minus the first clip's first ``skip`` frames, to ``out_path``. + + Returns ``(frame_counts, num_frames)`` for the WHOLE timeline — the skipped frames + are counted as part of the first clip and of the output — so the audio is laid out + the same whether or not the head was stream-copied. + """ + # Frames stream from the decoders straight into the encoder; only the + # transition windows are buffered. See _iter_joined_frames. + writer = make_mp4_writer(out_path, output_fps) + num_frames = 0 + frame_counts: list[int] = [] + decoders = [iter_video_frames(p, is_canceled=context.util.is_canceled) for p in paths] + try: + # Every clip goes through the canvas guard, including the first: frames already + # on it pass straight through, so nothing can reach the writer at a size the + # file is not being encoded at. + clip_iters = [ + _iter_frames_on_canvas(islice(frames, skip, None) if i == 0 else frames, canvas) + for i, frames in enumerate(decoders) + ] + for frame in self._iter_joined_frames( + clip_iters, is_canceled=context.util.is_canceled, frame_counts=frame_counts + ): + writer.append_data(frame) + num_frames += 1 finally: - for cleanup in (tmp_path, wav_path, muxed_path): - if cleanup is None: - continue - try: - cleanup.unlink(missing_ok=True) - except Exception: - pass + writer.close() + # An unfinished decoder holds the process-wide stream slot that the splice, or a + # fallback re-encode, would otherwise wait behind. + for frames in decoders: + frames.close() + frame_counts[0] += skip + return frame_counts, num_frames + skip + + def _splice( + self, + context: InvocationContext, + first_path: Path, + keyframe: int, + layout: H264Layout, + rest_path: Path, + first_clip_frames: int, + work_dir: Path, + out_path: Path, + ) -> bool: + """Join the first clip's frames before ``keyframe``, stream-copied, onto ``rest_path``. + + Returns False, leaving the caller to re-encode everything, when the two cannot be + joined losslessly: the remainder was encoded with a different decoder configuration + than the first clip carries (an upload, or a file from an older encoder setting), or + the first clip decoded to a different frame count than its sample table lists, which + would put the copied head and the skipped frames out of step. + """ + rest = read_h264_layout(rest_path) + if first_clip_frames != layout.sample_count: + reason = f"it decoded to {first_clip_frames} frames but lists {layout.sample_count}" + elif rest is None or rest.decoder_config != layout.decoder_config or rest.sync_samples[:1] != (0,): + reason = "its encoding differs from this node's" + else: + reason = None + if reason is None: + head_path = work_dir / "head.mp4" + try: + copy_frames(first_path, 0, keyframe, layout.sample_count, head_path) + head = read_h264_layout(head_path) + if head is None or head.sample_count != keyframe: + reason = "copying its head did not yield the expected frames" + else: + concat_stream_copy([head_path, rest_path], out_path) + joined = read_h264_layout(out_path) + if ( + joined is None + or joined.sample_count != keyframe + rest.sample_count + or joined.frame_duration != layout.frame_duration + ): + reason = "the joined file does not hold the expected frames on one timeline" + except SpliceError as e: + reason = str(e) + if reason is not None: + context.logger.info(f"Re-encoding the first clip in full: it cannot be stream-copied ({reason})") + return False + context.logger.info(f"Stream-copied the first {keyframe} frames of the first clip without re-encoding") + return True + + def _mux_spliced_audio( + self, + context: InvocationContext, + first_path: Path, + first_fps: Optional[float], + first_frames: int, + pcm: np.ndarray, + rate: int, + output_fps: float, + video_path: Path, + out_path: Path, + ) -> bool: + """Mux ``pcm`` onto ``video_path``, stream-copying the first clip's own AAC packets + for as much of the soundtrack's head as they already hold (see ``aac_splice``). + + The rebuilt soundtrack equals the first clip's decoded audio up to its outgoing + transition — provided that clip plays at the output rate, so its audio is not + retimed, and sets the soundtrack's sample rate. Returns False, leaving the caller to + encode the whole soundtrack, when the clip's track is not one this app's encoder + writes, nothing whole-frame precedes the transition, or the result's head does not + decode back to exactly the samples it should hold. + """ + if first_fps is None or first_fps != output_fps: + return False + layout, encoder = read_aac_layout(first_path), encoder_aac_layout(rate) + if layout is None or encoder is None or layout.sample_rate != rate: + return False + if (layout.decoder_config, layout.priming) != (encoder.decoder_config, encoder.priming): + return False + # The soundtrack holds the clip's own decoded samples only where its audio was mapped + # one-to-one: the clip's span at its native rate must be exactly its span on the timeline. + if round(first_frames / first_fps * rate) != round(first_frames * rate / output_fps): + return False + tail_need, _ = self._transition_windows() + unchanged = min( + round((first_frames - tail_need) * rate / output_fps), + layout.packet_count * AAC_FRAME_SAMPLES - layout.priming, + pcm.shape[1], + ) + preroll = PREROLL_FRAMES * AAC_FRAME_SAMPLES + # Keep the cut PREROLL_FRAMES clear of the first changed sample (see aac_splice). + head_packets = (unchanged + layout.priming) // AAC_FRAME_SAMPLES - PREROLL_FRAMES + cut = head_packets * AAC_FRAME_SAMPLES - layout.priming + if cut < preroll: + return False + try: + seam_reference = mux_spliced_aac( + video_path, + first_path, + head_packets, + layout.priming, + pcm[:, cut - preroll :], + rate, + pcm.shape[1], + out_path, + ) + muxed = extract_audio_pcm(out_path) + except (SpliceError, AudioExtractionError) as e: + context.logger.info(f"Re-encoding the whole soundtrack: it cannot be stream-copied ({e})") + return False + reason = self._spliced_audio_defect(muxed, pcm, rate, cut, seam_reference) + if reason is not None: + context.logger.info(f"Re-encoding the whole soundtrack: {reason}") + return False + context.logger.info(f"Stream-copied the first {cut} audio samples of the first clip without re-encoding") + return True + + @staticmethod + def _spliced_audio_defect( + muxed: Optional[tuple[np.ndarray, int]], + pcm: np.ndarray, + rate: int, + cut: int, + seam_reference: np.ndarray, + ) -> Optional[str]: + """Why the spliced soundtrack cannot be used, or None. + + The copied head must decode back to exactly the samples it stands for, and the + seam — the frames where copied and new packets overlap — must be no noisier than an + ordinary encode of the same samples (``seam_reference``). Whether the overlap + cancels depends on the encoder's window decisions around the content there, so it is + measured, not assumed; and it is measured against the seam's own content, since the + audio after it may be silence that codes with no error at all. + """ + if muxed is None or muxed[1] != rate or muxed[0].shape[1] < pcm.shape[1]: + return "the spliced soundtrack did not decode to its full length" + decoded = muxed[0] + if not np.array_equal(decoded[:, :cut], pcm[:, :cut]): + return "the copied audio did not decode back exactly" + seam = slice(cut, cut + seam_reference.shape[1]) + if seam_reference.shape[1] != PREROLL_FRAMES * AAC_FRAME_SAMPLES or seam.stop > pcm.shape[1]: + return "the new audio does not cover the seam" + seam_error = float(np.sqrt(np.mean((decoded[:, seam] - pcm[:, seam]) ** 2))) + coding_error = float(np.sqrt(np.mean((seam_reference - pcm[:, seam]) ** 2))) + if seam_error > 2.0 * coding_error + 1e-3: + return "the seam between copied and new audio is not clean" + return None def _estimate_transition_memory(self, width: int, height: int) -> int: if self.transition == "cut" or self.transition_frames == 0: @@ -346,14 +595,7 @@ def s(frames: int) -> int: """Frame count -> sample count on the emitted timeline.""" return round(frames * rate / output_fps) - # Boundary bookkeeping — must mirror _iter_joined_frames. - if self.transition == "crossfade" and self.transition_frames > 0: - tail_need = head_need = self.transition_frames - elif self.transition == "fade_through_black" and self.transition_frames > 0: - tail_need = self.transition_frames // 2 - head_need = self.transition_frames - tail_need - else: - tail_need = head_need = 0 + tail_need, head_need = self._transition_windows() clips: list[np.ndarray] = [] for i, (item, n_frames, native_fps) in enumerate(zip(extracted, frame_counts, native_rates, strict=True)): @@ -444,16 +686,10 @@ def _iter_joined_frames( transition windows. Because frames stream straight into the encoder, that error can surface mid-encode; the caller discards the partial output file. """ - if self.transition == "crossfade" and self.transition_frames > 0: - tail_need = head_need = self.transition_frames - blend = _crossfade - elif self.transition == "fade_through_black" and self.transition_frames > 0: - tail_need = self.transition_frames // 2 - head_need = self.transition_frames - tail_need - blend = _fade_through_black - else: + tail_need, head_need = self._transition_windows() + blend = {"crossfade": _crossfade, "fade_through_black": _fade_through_black}.get(self.transition) + if head_need + tail_need == 0: # "cut", or a zero-length transition: plain concatenation. - tail_need = head_need = 0 blend = None a_tail: list[np.ndarray] = [] diff --git a/invokeai/app/invocations/video_frame_extract_range.py b/invokeai/app/invocations/video_frame_extract_range.py index 0104cbc08ae..eb67c19ddd3 100644 --- a/invokeai/app/invocations/video_frame_extract_range.py +++ b/invokeai/app/invocations/video_frame_extract_range.py @@ -6,6 +6,13 @@ Concatenate Videos to splice clips together — e.g. trim a generated clip to a usable middle section before chaining it to another shot. +Frames are re-encoded only where they have to be. A range covering the whole +source at its own frame rate is a copy of the source file, so passing a clip +through this node (as video extension does with an untrimmed source) costs no +quality. A trim of a source this app encoded stream-copies the keyframe-aligned +run of frames inside the range and encodes only the partial keyframe intervals +at either end (see ``h264_splice``). Anything else is re-encoded in full. + If the source carries an audio track, the matching slice of it is carried through (and retimed alongside the video when the output fps differs from the source's). Like ``video_concat``, the audio path buffers the source @@ -14,7 +21,9 @@ rather than emit a silently-stripped clip. """ +import math import tempfile +from contextlib import closing from itertools import islice from pathlib import Path from typing import Callable, Iterator, Optional, Protocol @@ -39,8 +48,17 @@ from invokeai.app.invocations.primitives import VideoOutput from invokeai.app.services.session_processor.session_processor_common import CanceledException from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.app.util.h264_splice import ( + MIN_ENCODED_FRAMES, + H264Layout, + SpliceError, + concat_stream_copy, + copy_frames, + read_h264_layout, +) +from invokeai.app.util.mp4_metadata import write_mp4_tags from invokeai.app.util.video_audio import extract_audio_pcm, mux_audio_into_video, resample_linear -from invokeai.app.util.video_encoding import make_mp4_writer, write_stereo_wav +from invokeai.app.util.video_encoding import make_mp4_writer, write_stereo_wav, writer_decoder_config from invokeai.app.util.video_thumbnails import decoder_frame_count, iter_video_frames, probe_video @@ -161,7 +179,8 @@ def invoke(self, context: InvocationContext) -> ExtractVideoRangeOutput: video_path = context.videos.get_path(self.video.video_name) width, height, duration, source_fps = probe_video(video_path) - n_frames = decoder_frame_count(video_path) + exact_frame_count = decoder_frame_count(video_path) + n_frames = exact_frame_count if n_frames is None: if not source_fps or duration <= 0: raise ValueError( @@ -195,6 +214,18 @@ def invoke(self, context: InvocationContext) -> ExtractVideoRangeOutput: # clear message instead of surfacing an opaque ffmpeg error mid-encode. _validate_even_dimensions(width, height, self.video.video_name) + # Frames can be copied rather than re-encoded only at the source's own rate (a stream copy + # cannot retime), and only from a track whose sample indices are the decoder's frame indices: + # read_h264_layout refuses edit-list trims and rotations, and the count must agree. + layout = None + if source_fps is not None and math.isclose(output_fps, source_fps, rel_tol=1e-6): + layout = read_h264_layout(video_path) + if layout is not None and layout.sample_count != exact_frame_count: + layout = None + # The whole clip is the source itself; copying it avoids another lossy generation. + if layout is not None and start == 0 and end == layout.sample_count - 1: + return self._save_unchanged(context, video_path, width, height, n_frames, output_fps) + context.util.signal_progress(f"Transcoding frames {start}-{end} of {n_frames} @ {output_fps:.2f} fps") tmp = tempfile.NamedTemporaryFile(prefix="invokeai_video_range_", suffix=".mp4", delete=False) @@ -203,19 +234,26 @@ def invoke(self, context: InvocationContext) -> ExtractVideoRangeOutput: wav_path: Optional[Path] = None muxed_path: Optional[Path] = None try: - # imageio's iter_index isn't exposed by iio.imiter, so we enumerate and skip. - # Frames stream straight from the decoder into the encoder; see _write_frame_range. - writer = make_mp4_writer(tmp_path, output_fps) - try: - num_frames = _write_frame_range( - iter_video_frames(video_path, is_canceled=context.util.is_canceled), - writer, - start, - end, - is_canceled=context.util.is_canceled, + num_frames: Optional[int] = None + copy_span = self._plan_copy_span(layout, start, end, width, height, output_fps) if layout else None + if copy_span is not None: + num_frames = self._write_spliced_range( + context, video_path, start, end, *copy_span, output_fps, tmp_path ) - finally: - writer.close() + if num_frames is None: + # imageio's iter_index isn't exposed by iio.imiter, so we enumerate and skip. + # Frames stream straight from the decoder into the encoder; see _write_frame_range. + writer = make_mp4_writer(tmp_path, output_fps) + try: + num_frames = _write_frame_range( + iter_video_frames(video_path, is_canceled=context.util.is_canceled), + writer, + start, + end, + is_canceled=context.util.is_canceled, + ) + finally: + writer.close() expected_frames = end - start + 1 if num_frames != expected_frames: @@ -286,6 +324,139 @@ def invoke(self, context: InvocationContext) -> ExtractVideoRangeOutput: except Exception: pass + @staticmethod + def _plan_copy_span( + layout: H264Layout, start: int, end: int, width: int, height: int, fps: float + ) -> Optional[tuple[int, int, H264Layout]]: + """The keyframe-aligned frames ``[k1, k2)`` inside ``[start, end]`` that can be + stream-copied rather than re-encoded, with the source's layout; None if there are none. + + ``k1`` is the first keyframe at or after ``start``, ``k2`` the last keyframe at or + before ``end + 1`` — or the end of the track when the range runs to the last frame — + each moved one keyframe inward when it would leave an encoded end shorter than + ``MIN_ENCODED_FRAMES``. The source must also carry the decoder configuration this + node's encoder writes and a constant frame duration, or the re-encoded ends could not + share its track. + """ + if layout.frame_duration is None or layout.decoder_config != writer_decoder_config(width, height, fps): + return None + boundaries = [*layout.sync_samples, layout.sample_count] + inner_starts = [k for k in layout.sync_samples if k >= start and not 0 < k - start < MIN_ENCODED_FRAMES] + inner_ends = [k for k in boundaries if k <= end + 1 and not 0 < end + 1 - k < MIN_ENCODED_FRAMES] + if not inner_starts or not inner_ends or inner_starts[0] >= inner_ends[-1]: + return None + return inner_starts[0], inner_ends[-1], layout + + def _write_spliced_range( + self, + context: InvocationContext, + video_path: Path, + start: int, + end: int, + k1: int, + k2: int, + layout: H264Layout, + output_fps: float, + out_path: Path, + ) -> Optional[int]: + """Write frames ``[start, end]`` to ``out_path``, stream-copying ``[k1, k2)``. + + Only the partial GOPs on either side — ``[start, k1)`` and ``[k2, end]`` — are + encoded. Returns the number of frames written, which the caller checks against the + range like a full encode's, or None (leaving the caller to re-encode the whole range) + when an encoded part's decoder configuration differs from the source's, so the parts + could not share one track. + """ + with tempfile.TemporaryDirectory(prefix="invokeai_video_range_") as work: + work_dir = Path(work) + head = (start, k1, work_dir / "head.mp4") + tail = (k2, end + 1, work_dir / "tail.mp4") + encoded = [span for span in (head, tail) if span[0] < span[1]] + writers = [make_mp4_writer(path, output_fps) for _, _, path in encoded] + written = 0 + # Closed explicitly: an open decoder holds the process-wide stream slot, which the + # copy and join below (and a fallback re-encode) must not wait behind. + with closing(iter_video_frames(video_path, is_canceled=context.util.is_canceled)) as frames: + try: + for idx, frame in enumerate(islice(frames, end + 1)): + if context.util.is_canceled(): + raise CanceledException + for (lo, hi, _), writer in zip(encoded, writers, strict=True): + if lo <= idx < hi: + writer.append_data(np.ascontiguousarray(frame)) + written += 1 + finally: + for writer in writers: + writer.close() + + middle_path = work_dir / "middle.mp4" + reason: Optional[str] = None + if any( + (part := read_h264_layout(path)) is None or part.decoder_config != layout.decoder_config + for _, _, path in encoded + ): + reason = "its encoding differs from this node's" + else: + try: + copy_frames(video_path, k1, k2, layout.sample_count, middle_path) + middle = read_h264_layout(middle_path) + if middle is None or middle.sample_count != k2 - k1: + reason = "copying its frames did not yield the expected count" + else: + parts = [head[2]] if head in encoded else [] + parts += [middle_path] + ([tail[2]] if tail in encoded else []) + concat_stream_copy(parts, out_path) + joined = read_h264_layout(out_path) + if ( + joined is None + or joined.sample_count != written + (k2 - k1) + or joined.frame_duration != layout.frame_duration + ): + reason = "the joined file does not hold the expected frames on one timeline" + except SpliceError as e: + reason = str(e) + if reason is not None: + context.logger.info(f"Re-encoding the whole range: the source cannot be stream-copied ({reason})") + return None + context.logger.info(f"Stream-copied frames {k1}-{k2 - 1} of the source without re-encoding") + return written + (k2 - k1) + + def _save_unchanged( + self, + context: InvocationContext, + video_path: Path, + width: int, + height: int, + n_frames: int, + fps: float, + ) -> ExtractVideoRangeOutput: + """Save a copy of the source as the output: the requested range is all of it.""" + with tempfile.TemporaryDirectory(prefix="invokeai_video_range_") as work: + copy_path = Path(work) / "source.mp4" + # A stream-copy remux that drops the source's keyed metadata: save() writes this + # node's record, and if that rewrite fails the source's must not be read back as it. + write_mp4_tags(video_path, copy_path, {}) + video_dto = context.videos.save( + source_path=copy_path, + width=width, + height=height, + duration=n_frames / fps, + fps=fps, + ) + context.logger.info(f"Saved {video_dto.video_name}: the whole of the source, copied without re-encoding") + base = VideoOutput.build(video_dto) + return ExtractVideoRangeOutput( + video=base.video, + source_video=self.video, + width=base.width, + height=base.height, + num_frames=base.num_frames, + fps=base.fps, + duration=base.duration, + start_frame=0, + end_frame=n_frames - 1, + ) + def _build_audio_track( self, context: InvocationContext, diff --git a/invokeai/app/util/aac_splice.py b/invokeai/app/util/aac_splice.py new file mode 100644 index 00000000000..e1458635804 --- /dev/null +++ b/invokeai/app/util/aac_splice.py @@ -0,0 +1,359 @@ +"""Lossless splicing of an MP4's AAC soundtrack by stream copy. + +The audio counterpart of ``h264_splice``. Re-encoding an AAC track costs quality on every +pass just as re-encoding video does, so a soundtrack rebuilt for a join keeps the leading +packets of the clip it starts with and encodes only what follows. + +AAC packets are not independent: each decodes to the overlap-add of its own inverse MDCT +with the second half of the previous packet's. Copying packets ``[0, m)`` therefore +reproduces the source's first ``m`` output frames exactly, and the first freshly encoded +packet after them blends with the last copied one. That blend cancels cleanly only when +both packets were coded from the same signal across their shared window, so two things +keep the seam at coding-noise level rather than a glitch: the new encode starts +``PREROLL_FRAMES`` frames before the cut (one is not enough — the encoder's first frames +after a cold start make different window decisions), and callers keep the cut that many +frames clear of anything the joined soundtrack changes. Callers also verify the seam. + +Packets are joined as ADTS, which carries no timing, so two things the MP4 encodes in +its edit list are restored afterwards: the encoder priming the source's first packet +holds (an ``-itsoffset`` when muxing, from which ffmpeg writes the leading edit), and the +exact end of the soundtrack (patched into the edit list, as ffmpeg's own encodes trim +the last packet's padding). + +Only a track this app's encoder could have written is spliced — the same +``AudioSpecificConfig`` and priming, constant 1024-sample packets — and callers verify +the copied head decodes bit-exact before trusting the result. +""" + +import struct +import tempfile +from pathlib import Path +from typing import BinaryIO, NamedTuple, Optional + +import numpy as np + +from invokeai.app.util.h264_splice import SpliceError, run_ffmpeg +from invokeai.app.util.mp4_metadata import find_box, iter_boxes +from invokeai.app.util.video_audio import AAC_ENCODE_ARGS, extract_audio_pcm + +AAC_FRAME_SAMPLES = 1024 +# Frames of already-known audio the new encode starts with before the cut, and the clearance +# the cut keeps from the first changed sample. +PREROLL_FRAMES = 2 +_MAX_TOP_LEVEL_BOXES = 1024 +_MAX_PACKETS = 1 << 24 +# AudioSampleEntry fields before its child boxes (SampleEntry 8 + version-0 audio header 20). +_AUDIO_SAMPLE_ENTRY_HEADER = 28 +_ES_DESCRIPTOR, _DECODER_CONFIG_DESCRIPTOR, _DECODER_SPECIFIC_INFO = 3, 4, 5 +# syncExtensionType 0x2b7, extension object type 5 (SBR), sbrPresentFlag 0, as ffmpeg writes it +# after a two-byte AAC-LC config. +_NO_SBR_EXTENSION = b"\x56\xe5\x00" + + +class AacLayout(NamedTuple): + """What an audio splice needs to know about an MP4's first audio track.""" + + packet_count: int + sample_rate: int + """The media timescale, which for AAC is the sample rate.""" + priming: int + """Leading samples the edit list skips: the encoder delay held by the first packet.""" + decoder_config: bytes + """The ``AudioSpecificConfig`` (object type, sample rate, channels), without the + backward-compatible "no SBR" extension ffmpeg appends to its own encodes but cannot + write for a stream muxed from ADTS (both describe the same plain AAC-LC stream).""" + + +def _read_descriptor(data: bytes, position: int) -> Optional[tuple[int, int, int]]: + """``(tag, payload_start, payload_end)`` of the MPEG-4 descriptor at ``position``.""" + if position >= len(data): + return None + tag, length, position = data[position], 0, position + 1 + for _ in range(4): + if position >= len(data): + return None + byte = data[position] + position += 1 + length = (length << 7) | (byte & 0x7F) + if not byte & 0x80: + break + if position + length > len(data): + return None + return tag, position, position + length + + +def _audio_specific_config(esds: bytes) -> Optional[bytes]: + es = _read_descriptor(esds, 4) # after the FullBox version/flags + if es is None or es[0] != _ES_DESCRIPTOR: + return None + flags = esds[es[1] + 2] if es[1] + 2 < es[2] else 0 + position = es[1] + 3 + if flags & 0x80: # streamDependenceFlag + position += 2 + if flags & 0x40: # URL_Flag + position += 1 + (esds[position] if position < es[2] else 0) + if flags & 0x20: # OCRstreamFlag + position += 2 + config = _read_descriptor(esds, position) + if config is None or config[0] != _DECODER_CONFIG_DESCRIPTOR: + return None + specific = _read_descriptor(esds, config[1] + 13) + if specific is None or specific[0] != _DECODER_SPECIFIC_INFO: + return None + return esds[specific[1] : specific[2]] + + +def read_aac_layout(path: Path) -> Optional[AacLayout]: + """Return the first audio track's layout if it is spliceable AAC, else None. + + Spliceable means one ``mp4a`` sample entry, 1024-sample packets (the last may be + shorter), and an edit list that only skips priming. Only I/O errors propagate. + """ + with open(path, "rb") as fh: + fh.seek(0, 2) + moov = find_box(fh, 0, fh.tell(), b"moov", max_boxes=_MAX_TOP_LEVEL_BOXES) + if moov is None: + return None + for box_type, trak_start, trak_end in iter_boxes(fh, *moov): + if box_type != b"trak": + continue + mdia = find_box(fh, trak_start, trak_end, b"mdia") + hdlr = find_box(fh, *mdia, b"hdlr") if mdia is not None else None + if hdlr is None or hdlr[0] + 12 > hdlr[1]: + continue + fh.seek(hdlr[0] + 8) + if fh.read(4) != b"soun": + continue + return _read_audio_track(fh, (trak_start, trak_end), mdia) + return None + + +def _read_audio_track(fh: BinaryIO, trak: tuple[int, int], mdia: tuple[int, int]) -> Optional[AacLayout]: + mdhd = find_box(fh, *mdia, b"mdhd") + minf = find_box(fh, *mdia, b"minf") + stbl = find_box(fh, *minf, b"stbl") if minf is not None else None + if mdhd is None or stbl is None or mdhd[0] + 24 > mdhd[1]: + return None + fh.seek(mdhd[0]) + fh.seek(mdhd[0] + (20 if fh.read(1)[0] == 1 else 12)) + (sample_rate,) = struct.unpack(">I", fh.read(4)) + + stsd = find_box(fh, *stbl, b"stsd") + if stsd is None or stsd[0] + 8 > stsd[1]: + return None + fh.seek(stsd[0] + 4) + if struct.unpack(">I", fh.read(4)) != (1,): + return None + entries = list(iter_boxes(fh, stsd[0] + 8, stsd[1], max_boxes=1)) + if not entries or entries[0][0] != b"mp4a": + return None + _, entry_start, entry_end = entries[0] + fh.seek(entry_start + 8) + if fh.read(2) != b"\x00\x00": # QuickTime v1/v2 sound entries lay their fields out differently + return None + esds = find_box(fh, entry_start + _AUDIO_SAMPLE_ENTRY_HEADER, entry_end, b"esds") + if esds is None: + return None + fh.seek(esds[0]) + decoder_config = _audio_specific_config(fh.read(esds[1] - esds[0])) + if not decoder_config: + return None + if len(decoder_config) == 5 and decoder_config.endswith(_NO_SBR_EXTENSION): + decoder_config = decoder_config[:2] + + stts = find_box(fh, *stbl, b"stts") + if stts is None or stts[0] + 8 > stts[1]: + return None + fh.seek(stts[0] + 4) + (runs,) = struct.unpack(">I", fh.read(4)) + if not 1 <= runs <= 2 or stts[0] + 8 + 8 * runs > stts[1]: + return None + flat = struct.unpack(f">{2 * runs}I", fh.read(8 * runs)) + counts, deltas = flat[::2], flat[1::2] + # Every packet holds one AAC frame; ffmpeg shortens only the last one, to its real samples. + if deltas[0] != AAC_FRAME_SAMPLES or (runs == 2 and (counts[1] != 1 or deltas[1] > AAC_FRAME_SAMPLES)): + return None + packet_count = sum(counts) + if packet_count > _MAX_PACKETS: + return None + + edts = find_box(fh, *trak, b"edts") + elst = find_box(fh, *edts, b"elst") if edts is not None else None + priming = 0 + if elst is not None: + if elst[0] + 8 > elst[1]: + return None + fh.seek(elst[0]) + version = fh.read(1)[0] + entry_format = ">QqhH" if version == 1 else ">IihH" + fh.seek(elst[0] + 4) + if struct.unpack(">I", fh.read(4)) != (1,) or elst[0] + 8 + struct.calcsize(entry_format) > elst[1]: + return None + _, priming, rate_integer, rate_fraction = struct.unpack(entry_format, fh.read(struct.calcsize(entry_format))) + if (rate_integer, rate_fraction) != (1, 0) or priming < 0: + return None + return AacLayout(packet_count, sample_rate, priming, decoder_config) + + +_encoder_layouts: dict[int, AacLayout] = {} + + +def encoder_aac_layout(sample_rate: int) -> Optional[AacLayout]: + """The layout ``mux_audio_into_video``'s encoder writes at ``sample_rate`` (stereo), or + None if the probe fails. Only a successful probe is remembered.""" + if sample_rate not in _encoder_layouts: + layout = _probe_encoder_layout(sample_rate) + if layout is None: + return None + _encoder_layouts[sample_rate] = layout + return _encoder_layouts[sample_rate] + + +def _probe_encoder_layout(sample_rate: int) -> Optional[AacLayout]: + from invokeai.app.util.video_encoding import write_stereo_wav + + with tempfile.TemporaryDirectory(prefix="invokeai_aac_probe_") as work: + wav, probe = Path(work) / "probe.wav", Path(work) / "probe.m4a" + write_stereo_wav(wav, np.zeros((2, 4 * AAC_FRAME_SAMPLES), dtype=np.float32), sample_rate) + try: + run_ffmpeg(["-i", str(wav), *AAC_ENCODE_ARGS, str(probe)], "probing the AAC encoder") + except SpliceError: + return None + return read_aac_layout(probe) + + +def _adts_frames(data: bytes) -> list[bytes]: + frames: list[bytes] = [] + position = 0 + while position < len(data): + header = data[position : position + 7] + if len(header) < 7 or header[0] != 0xFF or header[1] & 0xF6 != 0xF0: + raise SpliceError("ffmpeg wrote a malformed ADTS stream") + length = ((header[3] & 0x03) << 11) | (header[4] << 3) | (header[5] >> 5) + if length < 7 or position + length > len(data): + raise SpliceError("ffmpeg wrote a truncated ADTS frame") + frames.append(data[position : position + length]) + position += length + return frames + + +def _set_track_duration(path: Path, handler: bytes, samples: int, sample_rate: int) -> None: + """Trim the ``handler`` track's presentation to ``samples`` via its edit list, and keep + ``tkhd`` and ``mvhd`` durations consistent with it. All fields are rewritten in place.""" + with open(path, "r+b") as fh: + fh.seek(0, 2) + moov = find_box(fh, 0, fh.tell(), b"moov", max_boxes=_MAX_TOP_LEVEL_BOXES) + mvhd = find_box(fh, *moov, b"mvhd") if moov is not None else None + if mvhd is None: + raise SpliceError(f"{path.name} has no movie header") + fh.seek(mvhd[0]) + mvhd_version = fh.read(1)[0] + fh.seek(mvhd[0] + (20 if mvhd_version == 1 else 12)) + (movie_timescale,) = struct.unpack(">I", fh.read(4)) + duration = round(samples * movie_timescale / sample_rate) + track_durations: list[int] = [] + patched = False + for box_type, trak_start, trak_end in iter_boxes(fh, *moov): + if box_type != b"trak": + continue + tkhd = find_box(fh, trak_start, trak_end, b"tkhd") + mdia = find_box(fh, trak_start, trak_end, b"mdia") + hdlr = find_box(fh, *mdia, b"hdlr") if mdia is not None else None + if tkhd is None or hdlr is None: + raise SpliceError(f"{path.name} has a malformed track") + fh.seek(hdlr[0] + 8) + is_target = fh.read(4) == handler + fh.seek(tkhd[0]) + wide = fh.read(1)[0] == 1 + tkhd_duration_at, field = tkhd[0] + (28 if wide else 20), ">Q" if wide else ">I" + if is_target and not patched: + edts = find_box(fh, trak_start, trak_end, b"edts") + elst = find_box(fh, *edts, b"elst") if edts is not None else None + if elst is None: + raise SpliceError(f"{path.name}'s {handler.decode()} track has no edit list") + fh.seek(elst[0]) + elst_field = ">Q" if fh.read(1)[0] == 1 else ">I" + fh.seek(elst[0] + 8) + fh.write(struct.pack(elst_field, duration)) + fh.seek(tkhd_duration_at) + fh.write(struct.pack(field, duration)) + patched = True + fh.seek(tkhd_duration_at) + track_durations.append(struct.unpack(field, fh.read(struct.calcsize(field)))[0]) + if not patched: + raise SpliceError(f"{path.name} has no {handler.decode()} track") + fh.seek(mvhd[0] + (24 if mvhd_version == 1 else 16)) + fh.write(struct.pack(">Q" if mvhd_version == 1 else ">I", max(track_durations))) + + +def mux_spliced_aac( + video_path: Path, + source: Path, + head_packets: int, + priming: int, + tail_pcm: np.ndarray, + sample_rate: int, + total_samples: int, + out_path: Path, +) -> np.ndarray: + """Write ``out_path`` = ``video_path``'s video plus a soundtrack spliced from ``source``. + + The soundtrack is ``source``'s first ``head_packets`` AAC packets, stream-copied — + ``head_packets * 1024 - priming`` samples after priming — followed by ``tail_pcm`` + encoded from ``PREROLL_FRAMES`` frames *before* that cut (``tail_pcm`` must start + there), presented for ``total_samples`` samples in all. + + Returns the new encode's own decode of the ``PREROLL_FRAMES`` frames after the cut: + the coding noise this content gets from an ordinary encode, which is what a clean + seam over the same samples should be measured against. + """ + from invokeai.app.util.video_encoding import write_stereo_wav + + if priming % AAC_FRAME_SAMPLES: + raise SpliceError(f"priming of {priming} samples is not a whole AAC frame") + with tempfile.TemporaryDirectory(prefix="invokeai_aac_splice_") as work: + head, tail_wav, tail = Path(work) / "head.aac", Path(work) / "tail.wav", Path(work) / "tail.aac" + run_ffmpeg( + ["-i", str(source), "-map", "0:a:0", "-c", "copy", "-frames:a", str(head_packets), "-f", "adts", str(head)], + f"copying the first {head_packets} audio packets of {source.name}", + ) + write_stereo_wav(tail_wav, tail_pcm, sample_rate) + run_ffmpeg(["-i", str(tail_wav), *AAC_ENCODE_ARGS, "-f", "adts", str(tail)], "encoding the new audio") + head_frames = _adts_frames(head.read_bytes()) + if len(head_frames) != head_packets: + raise SpliceError(f"copied {len(head_frames)} audio packets, expected {head_packets}") + # Drop the new encode's priming packets and its pre-roll. + dropped = priming // AAC_FRAME_SAMPLES + PREROLL_FRAMES + tail_frames = _adts_frames(tail.read_bytes())[dropped:] + decoded = extract_audio_pcm(tail) + if decoded is None: + raise SpliceError("the new audio decoded to nothing") + # An ADTS decode knows nothing of priming, so it starts with the priming frame. + seam_start = priming + PREROLL_FRAMES * AAC_FRAME_SAMPLES + seam_reference = decoded[0][:, seam_start : seam_start + PREROLL_FRAMES * AAC_FRAME_SAMPLES] + joined = Path(work) / "joined.aac" + joined.write_bytes(b"".join(head_frames + tail_frames)) + run_ffmpeg( + [ + "-i", + str(video_path), + # The first copied packet still holds the source's priming; presenting it from + # -priming makes ffmpeg write the leading edit that skips it. + "-itsoffset", + repr(-priming / sample_rate), + "-i", + str(joined), + "-map", + "0:v:0", + "-map", + "1:a:0", + "-c", + "copy", + "-f", + "mp4", + str(out_path), + ], + "muxing the spliced audio", + ) + _set_track_duration(out_path, b"soun", total_samples, sample_rate) + return seam_reference diff --git a/invokeai/app/util/h264_splice.py b/invokeai/app/util/h264_splice.py new file mode 100644 index 00000000000..83af5c7459f --- /dev/null +++ b/invokeai/app/util/h264_splice.py @@ -0,0 +1,444 @@ +"""Lossless splicing of H.264 MP4s by stream copy. + +Re-encoding a clip that is already H.264 costs quality on every pass, and the losses +compound: a video extended N times by decode-everything-and-re-encode carries its first +segment through N generations of artifacts. When the head of a clip can be reused as-is, +:func:`copy_frames` cuts out a keyframe-aligned run of frames without decoding, and +:func:`concat_stream_copy` joins it to freshly encoded neighbours. + +That join is only valid when every part carries the *same* decoder configuration: an MP4 +stores one ``avcC`` (profile, level, SPS and PPS) per track, so a part encoded with +different parameters would be decoded against the wrong headers. :func:`read_h264_layout` +exposes the ``avcC`` bytes so callers can require an exact match. Parts written by this +app's own encoder (``make_mp4_writer``) at the same canvas and frame rate match; anything +else is reported as not spliceable rather than guessed at. + +Cuts are made only at IDR pictures. Nothing after an IDR in decode order references +anything before it, and everything before it is displayed before it, so the samples +between two IDRs decode to exactly the frames displayed between them. A plain sync +sample is not enough: an open-GOP I-frame is marked as one too, and the frames that +follow it in decode order but display before it would be lost from the copied run. + +A joined file also has to keep one timeline. x264 gives a part of fewer than +``MIN_ENCODED_FRAMES`` frames a shorter B-frame reorder delay than a full one, and the +concat demuxer then squeezes timestamps at the seam; callers keep encoded parts out of +that range and check :attr:`H264Layout.frame_duration` of the result. + +Sample indices only name displayed frames when nothing in the container reinterprets the +samples, so two more things make a track unspliceable: an edit list that hides samples (a +``-c copy`` trim keeps every sample and cuts with an edit list instead, so decoding yields +fewer frames than the sample table lists), and a display matrix other than the identity (a +rotation applied as metadata, which decoding applies to the frames this app re-encodes but +a stream copy carries into the joined file for all of them). +""" + +import shutil +import struct +import subprocess +import tempfile +from pathlib import Path +from typing import BinaryIO, NamedTuple, Optional + +import imageio_ffmpeg + +from invokeai.app.util.mp4_metadata import find_box, iter_boxes + +# A stream copy is disk-bound; this only stops a wedged child from holding the node forever. +SPLICE_TIMEOUT_SECONDS = 600 +# ``moov`` precedes or follows ``mdat`` in a real file; a dozen top-level boxes at most. +_MAX_TOP_LEVEL_BOXES = 1024 +# A sync-sample table for even an hour of 120 fps all-intra video is ~430k entries. +_MAX_SYNC_ENTRIES = 1 << 20 +# VisualSampleEntry fields before its child boxes (SampleEntry 8 + visual header 70). +_VISUAL_SAMPLE_ENTRY_HEADER = 78 +# Shortest part x264 (bframes=3, as make_mp4_writer leaves it) encodes with its full B-frame +# reorder delay; a shorter part joined to full-delay parts distorts the joined timestamps. +MIN_ENCODED_FRAMES = 3 +# A sample table for an hour of 120 fps video is ~430k entries. +_MAX_SAMPLES = 1 << 22 +# NAL unit types: 1-5 are coded slices, 5 an IDR picture's. +_NAL_IDR = 5 +_NAL_VCL = range(1, 6) +# The 3x3 ``tkhd`` display matrix that leaves frames as decoded (16.16 / 2.30 fixed point). +_IDENTITY_MATRIX = (0x10000, 0, 0, 0, 0x10000, 0, 0, 0, 0x40000000) + + +class H264Layout(NamedTuple): + """What a splice needs to know about an MP4's first video track.""" + + sample_count: int + """Number of video samples (frames) in the track.""" + sync_samples: tuple[int, ...] + """0-based decode-order indices of the IDR pictures — the only clean cut points — ascending.""" + decoder_config: bytes + """The raw ``avcC`` payload: profile, level, SPS and PPS.""" + frame_duration: Optional[int] + """Every sample's duration in media ticks, or None if they are not all equal.""" + + +class SpliceError(RuntimeError): + """ffmpeg failed to copy or join the streams.""" + + +def _read_full_box_table(fh: BinaryIO, start: int, end: int, header: str) -> Optional[tuple[int, ...]]: + """Read a FullBox's fixed header fields (after version/flags) as big-endian uint32s.""" + size = struct.calcsize(header) + if start + 4 + size > end: + return None + fh.seek(start + 4) + return struct.unpack(header, fh.read(size)) + + +def read_h264_layout(path: Path) -> Optional[H264Layout]: + """Return the first video track's layout if it is spliceable H.264, else None. + + None covers anything this module will not splice: no video track, a codec other than + ``avc1``, a fragmented file (its samples live in ``moof`` boxes, not the sample table), + an edit list that hides samples, a non-identity display matrix (see the module + docstring), and truncated or malformed boxes. Only I/O errors propagate. + """ + with open(path, "rb") as fh: + fh.seek(0, 2) + file_size = fh.tell() + moov = find_box(fh, 0, file_size, b"moov", max_boxes=_MAX_TOP_LEVEL_BOXES) + if moov is None: + return None + for box_type, trak_start, trak_end in iter_boxes(fh, *moov): + if box_type != b"trak": + continue + mdia = find_box(fh, trak_start, trak_end, b"mdia") + if mdia is None: + continue + hdlr = find_box(fh, *mdia, b"hdlr") + if hdlr is None or hdlr[0] + 12 > hdlr[1]: + continue + fh.seek(hdlr[0] + 8) + if fh.read(4) != b"vide": + continue + # First video track only: that is the stream every ffmpeg call here maps (0:v:0). + minf = find_box(fh, *mdia, b"minf") + stbl = find_box(fh, *minf, b"stbl") if minf is not None else None + movie_timescale = _read_timescale(fh, find_box(fh, *moov, b"mvhd")) + media_timescale = _read_timescale(fh, find_box(fh, *mdia, b"mdhd")) + if stbl is None or movie_timescale is None or media_timescale is None: + return None + if not _has_identity_matrix(fh, find_box(fh, trak_start, trak_end, b"tkhd")): + return None + edts = find_box(fh, trak_start, trak_end, b"edts") + if edts is not None and not _edits_show_every_sample( + fh, find_box(fh, *edts, b"elst"), stbl, movie_timescale, media_timescale + ): + return None + return _read_video_stbl(fh, *stbl) + return None + + +def _read_timescale(fh: BinaryIO, box: Optional[tuple[int, int]]) -> Optional[int]: + """The timescale of an ``mvhd`` or ``mdhd`` box (same layout up to that field).""" + if box is None or box[0] + 24 > box[1]: + return None + fh.seek(box[0]) + offset = 20 if fh.read(1)[0] == 1 else 12 + fh.seek(box[0] + offset) + (timescale,) = struct.unpack(">I", fh.read(4)) + return timescale or None + + +def _has_identity_matrix(fh: BinaryIO, tkhd: Optional[tuple[int, int]]) -> bool: + if tkhd is None or tkhd[0] >= tkhd[1]: + return False + fh.seek(tkhd[0]) + offset = 52 if fh.read(1)[0] == 1 else 40 + if tkhd[0] + offset + 36 > tkhd[1]: + return False + fh.seek(tkhd[0] + offset) + return struct.unpack(">9i", fh.read(36)) == _IDENTITY_MATRIX + + +def _edits_show_every_sample( + fh: BinaryIO, + elst: Optional[tuple[int, int]], + stbl: tuple[int, int], + movie_timescale: int, + media_timescale: int, +) -> bool: + """True if the edit list only shifts presentation to the first frame, as ffmpeg writes + for B-frame reordering delay, rather than hiding any samples. + + That edit is a single normal-rate entry starting at the first sample's composition + offset and lasting the whole track. A trim (later start or shorter duration), an empty + leading edit, or several edits all change which samples are displayed. + """ + if elst is None or elst[0] + 8 > elst[1]: + return False + fh.seek(elst[0]) + version = fh.read(1)[0] + entry_format = ">QqhH" if version == 1 else ">IihH" + entry_size = struct.calcsize(entry_format) + fh.seek(elst[0] + 4) + (entry_count,) = struct.unpack(">I", fh.read(4)) + if entry_count != 1 or elst[0] + 8 + entry_size > elst[1]: + return False + segment_duration, media_time, rate_integer, rate_fraction = struct.unpack(entry_format, fh.read(entry_size)) + if (rate_integer, rate_fraction) != (1, 0): + return False + + runs = _read_stts(fh, stbl) + if runs is None: + return False + track_duration = sum(count * delta for count, delta in runs) + first_delta = runs[0][1] + + first_offset = 0 + ctts = find_box(fh, *stbl, b"ctts") + if ctts is not None: + if ctts[0] + 16 > ctts[1]: + return False + fh.seek(ctts[0] + 12) + (first_offset,) = struct.unpack(">i", fh.read(4)) + if media_time != first_offset: + return False + # The segment is in movie units, usually rounded to milliseconds: allow that rounding, + # but not a missing frame. + return segment_duration * media_timescale >= (track_duration - first_delta // 2) * movie_timescale + + +def _read_stts(fh: BinaryIO, stbl: tuple[int, int]) -> Optional[list[tuple[int, int]]]: + """The ``(sample_count, sample_delta)`` runs of the time-to-sample table.""" + stts = find_box(fh, *stbl, b"stts") + if stts is None or stts[0] + 8 > stts[1]: + return None + fh.seek(stts[0] + 4) + (count,) = struct.unpack(">I", fh.read(4)) + if count == 0 or count > _MAX_SAMPLES or stts[0] + 8 + 8 * count > stts[1]: + return None + flat = struct.unpack(f">{2 * count}I", fh.read(8 * count)) + return list(zip(flat[::2], flat[1::2], strict=True)) + + +def _sample_locations(fh: BinaryIO, stbl: tuple[int, int], sample_count: int) -> Optional[list[tuple[int, int]]]: + """``(file_offset, size)`` of every sample, from the size, chunk and chunk-offset tables.""" + stsz, stsc = find_box(fh, *stbl, b"stsz"), find_box(fh, *stbl, b"stsc") + stco, co64 = find_box(fh, *stbl, b"stco"), find_box(fh, *stbl, b"co64") + offsets_box, offset_format = (stco, "I") if stco is not None else (co64, "Q") + if stsz is None or stsc is None or offsets_box is None or sample_count > _MAX_SAMPLES: + return None + # Each table opens with version/flags and a 4-byte count (stsz: a size and a count). + if any(box[0] + 8 > box[1] for box in (stsz, stsc, offsets_box)): + return None + fh.seek(stsz[0] + 4) + (uniform_size,) = struct.unpack(">I", fh.read(4)) + if uniform_size: + sizes: tuple[int, ...] = (uniform_size,) * sample_count + else: + if stsz[0] + 12 + 4 * sample_count > stsz[1]: + return None + fh.seek(stsz[0] + 12) + sizes = struct.unpack(f">{sample_count}I", fh.read(4 * sample_count)) + fh.seek(stsc[0] + 4) + (run_count,) = struct.unpack(">I", fh.read(4)) + if run_count == 0 or stsc[0] + 8 + 12 * run_count > stsc[1]: + return None + flat = struct.unpack(f">{3 * run_count}I", fh.read(12 * run_count)) + runs = [(flat[i], flat[i + 1]) for i in range(0, len(flat), 3)] + fh.seek(offsets_box[0] + 4) + (chunk_count,) = struct.unpack(">I", fh.read(4)) + entry = struct.calcsize(offset_format) + if chunk_count > _MAX_SAMPLES or offsets_box[0] + 8 + entry * chunk_count > offsets_box[1]: + return None + chunk_offsets = struct.unpack(f">{chunk_count}{offset_format}", fh.read(entry * chunk_count)) + + locations: list[tuple[int, int]] = [] + for run, (first_chunk, per_chunk) in enumerate(runs): + last_chunk = runs[run + 1][0] - 1 if run + 1 < len(runs) else chunk_count + if not 1 <= first_chunk <= last_chunk + 1 or last_chunk > chunk_count: + return None + for chunk in range(first_chunk - 1, last_chunk): + position = chunk_offsets[chunk] + for _ in range(per_chunk): + if len(locations) == sample_count: + return None + size = sizes[len(locations)] + locations.append((position, size)) + position += size + return locations if len(locations) == sample_count else None + + +def _is_idr(fh: BinaryIO, offset: int, size: int, length_size: int) -> bool: + """True if the sample's first coded slice is an IDR slice (AUD, SEI, SPS and PPS may precede it).""" + position, end = offset, offset + size + while position + length_size + 1 <= end: + fh.seek(position) + header = fh.read(length_size + 1) + if len(header) < length_size + 1: + return False + nal_length = int.from_bytes(header[:length_size], "big") + nal_type = header[length_size] & 0x1F + if nal_type in _NAL_VCL: + return nal_type == _NAL_IDR + if nal_length == 0: + return False + position += length_size + nal_length + return False + + +def _read_video_stbl(fh: BinaryIO, stbl_start: int, stbl_end: int) -> Optional[H264Layout]: + stbl = (stbl_start, stbl_end) + stsd = find_box(fh, *stbl, b"stsd") + stsz = find_box(fh, *stbl, b"stsz") + if stsd is None or stsz is None: + return None + + # stsd: version/flags, entry_count, then sample entries. One entry is all ffmpeg writes; + # several would mean samples switch configuration mid-track, which a splice cannot match. + entry_count = _read_full_box_table(fh, *stsd, ">I") + if entry_count != (1,): + return None + entries = list(iter_boxes(fh, stsd[0] + 8, stsd[1], max_boxes=1)) + if not entries or entries[0][0] != b"avc1": + return None + _, entry_start, entry_end = entries[0] + avcc = find_box(fh, entry_start + _VISUAL_SAMPLE_ENTRY_HEADER, entry_end, b"avcC") + if avcc is None or avcc[0] >= avcc[1]: + return None + fh.seek(avcc[0]) + decoder_config = fh.read(avcc[1] - avcc[0]) + + sizes = _read_full_box_table(fh, *stsz, ">II") + if sizes is None: + return None + sample_count = sizes[1] + if sample_count == 0: + return None + + stss = find_box(fh, *stbl, b"stss") + if stss is None: + # No sync-sample table means every sample is a sync sample. + sync_samples = tuple(range(sample_count)) + else: + count = _read_full_box_table(fh, *stss, ">I") + if count is None or count[0] > _MAX_SYNC_ENTRIES or stss[0] + 8 + 4 * count[0] > stss[1]: + return None + fh.seek(stss[0] + 8) + one_based = struct.unpack(f">{count[0]}I", fh.read(4 * count[0])) + if any(not 1 <= n <= sample_count for n in one_based) or list(one_based) != sorted(set(one_based)): + return None + sync_samples = tuple(n - 1 for n in one_based) + + if len(decoder_config) < 5: + return None + length_size = (decoder_config[4] & 0x03) + 1 + locations = _sample_locations(fh, stbl, sample_count) + if locations is None: + return None + idr_samples = tuple(k for k in sync_samples if _is_idr(fh, *locations[k], length_size)) + + runs = _read_stts(fh, stbl) + if runs is None: + return None + deltas = {delta for _, delta in runs} + frame_duration = deltas.pop() if len(deltas) == 1 else None + return H264Layout( + sample_count=sample_count, + sync_samples=idr_samples, + decoder_config=decoder_config, + frame_duration=frame_duration, + ) + + +def run_ffmpeg(args: list[str], description: str) -> None: + try: + proc = subprocess.run( + [imageio_ffmpeg.get_ffmpeg_exe(), "-y", "-loglevel", "error", *args], + capture_output=True, + timeout=SPLICE_TIMEOUT_SECONDS, + ) + except subprocess.TimeoutExpired as e: + raise SpliceError(f"ffmpeg timed out {description}") from e + if proc.returncode != 0: + stderr = proc.stderr.decode("utf-8", errors="replace").strip() + raise SpliceError(f"ffmpeg failed {description}: {stderr[-500:]}") + + +def copy_frames(src: Path, start: int, end: int, sample_count: int, dst: Path) -> None: + """Write ``dst`` = frames ``[start, end)`` of ``src``'s video, stream-copied. + + ``start`` and ``end`` are decode-order sample indices that must each be a keyframe (or + 0, or ``sample_count`` for the end of the track) so the cuts are clean: with closed + GOPs, the packets between two keyframes decode to exactly the frames displayed between + them. Audio is dropped; callers rebuild the soundtrack for the new timeline. + """ + description = f"copying frames {start}-{end - 1} of {src.name}" + if start == 0: + run_ffmpeg( + ["-i", str(src), "-map", "0:v:0", "-c", "copy", "-frames:v", str(end), "-f", "mp4", str(dst)], + description, + ) + return + # The segment muxer splits at exact packet counts; the wanted range is the second segment. + split_points = [start] + ([end] if end < sample_count else []) + with tempfile.TemporaryDirectory(prefix="invokeai_h264_splice_") as work: + pattern = Path(work) / "segment%d.mp4" + run_ffmpeg( + [ + "-i", + str(src), + "-map", + "0:v:0", + "-c", + "copy", + "-f", + "segment", + "-segment_format", + "mp4", + "-segment_frames", + ",".join(str(n) for n in split_points), + "-reset_timestamps", + "1", + str(pattern), + ], + description, + ) + segment = Path(work) / "segment1.mp4" + if not segment.exists(): + raise SpliceError(f"ffmpeg wrote no segment {description}") + shutil.move(str(segment), str(dst)) + + +def _concat_list_entry(path: Path) -> str: + # The concat demuxer's quoting: close the quote, emit an escaped quote, reopen. + escaped = str(path.resolve()).replace("'", "'\\''") + return f"file '{escaped}'\n" + + +def concat_stream_copy(parts: list[Path], dst: Path) -> None: + """Write ``dst`` = ``parts``' video streams joined end to end, stream-copied. + + Every part must share one decoder configuration (see :func:`read_h264_layout`) and + start on a keyframe; the concat demuxer offsets each part's timestamps by the duration + of the parts before it. + """ + with tempfile.NamedTemporaryFile("w", encoding="utf-8", suffix=".txt", delete=False) as listing: + listing.write("ffconcat version 1.0\n" + "".join(_concat_list_entry(p) for p in parts)) + list_path = Path(listing.name) + try: + run_ffmpeg( + [ + "-f", + "concat", + "-safe", + "0", + "-i", + str(list_path), + "-map", + "0:v:0", + "-c", + "copy", + "-f", + "mp4", + str(dst), + ], + f"joining {len(parts)} stream-copied parts", + ) + finally: + list_path.unlink(missing_ok=True) diff --git a/invokeai/app/util/mp4_metadata.py b/invokeai/app/util/mp4_metadata.py index 1b1b6e32d49..d4b1a9759d9 100644 --- a/invokeai/app/util/mp4_metadata.py +++ b/invokeai/app/util/mp4_metadata.py @@ -56,9 +56,7 @@ def _ffmpeg_exe() -> str: return imageio_ffmpeg.get_ffmpeg_exe() -def _iter_boxes( - fh: BinaryIO, start: int, end: int, max_boxes: Optional[int] = None -) -> Iterator[tuple[bytes, int, int]]: +def iter_boxes(fh: BinaryIO, start: int, end: int, max_boxes: Optional[int] = None) -> Iterator[tuple[bytes, int, int]]: """Yield ``(type, payload_start, payload_end)`` for each box in ``[start, end)``. Stops silently at the first structurally impossible box (size smaller than its own @@ -91,10 +89,10 @@ def _iter_boxes( position += size -def _find_box( +def find_box( fh: BinaryIO, start: int, end: int, box_type: bytes, max_boxes: Optional[int] = None ) -> Optional[tuple[int, int]]: - for found_type, payload_start, payload_end in _iter_boxes(fh, start, end, max_boxes): + for found_type, payload_start, payload_end in iter_boxes(fh, start, end, max_boxes): if found_type == box_type: return payload_start, payload_end return None @@ -141,7 +139,7 @@ def _read_keys(fh: BinaryIO, start: int, end: int) -> dict[int, str]: def _read_utf8_data(fh: BinaryIO, start: int, end: int) -> Optional[str]: """The first UTF-8 ``data`` atom of an ``ilst`` item, or None.""" - for box_type, payload_start, payload_end in _iter_boxes(fh, start, end): + for box_type, payload_start, payload_end in iter_boxes(fh, start, end): if box_type != b"data" or payload_end - payload_start < 8: continue fh.seek(payload_start) @@ -169,22 +167,22 @@ def read_mp4_tags(path: Path, keys: Optional[Collection[str]] = None) -> dict[st with open(path, "rb") as fh: fh.seek(0, 2) file_size = fh.tell() - moov = _find_box(fh, 0, file_size, b"moov", max_boxes=_MAX_TOP_LEVEL_BOXES) + moov = find_box(fh, 0, file_size, b"moov", max_boxes=_MAX_TOP_LEVEL_BOXES) if moov is None: return result - udta = _find_box(fh, *moov, b"udta") + udta = find_box(fh, *moov, b"udta") if udta is None: return result - meta = _find_box(fh, *udta, b"meta") + meta = find_box(fh, *udta, b"meta") if meta is None: return result meta_start = _meta_payload_start(fh, *meta) - keys_box = _find_box(fh, meta_start, meta[1], b"keys") - ilst = _find_box(fh, meta_start, meta[1], b"ilst") + keys_box = find_box(fh, meta_start, meta[1], b"keys") + ilst = find_box(fh, meta_start, meta[1], b"ilst") if keys_box is None or ilst is None: return result names = _read_keys(fh, *keys_box) - for item_type, item_start, item_end in _iter_boxes(fh, *ilst): + for item_type, item_start, item_end in iter_boxes(fh, *ilst): name = names.get(int.from_bytes(item_type, "big")) if name is None or (wanted is not None and name not in wanted): continue @@ -200,7 +198,7 @@ def read_ftyp_major_brand(path: Path, search_limit: int = 64 * 1024) -> Optional with open(path, "rb") as fh: fh.seek(0, 2) end = min(fh.tell(), search_limit) - for box_type, payload_start, payload_end in _iter_boxes(fh, 0, end): + for box_type, payload_start, payload_end in iter_boxes(fh, 0, end): if box_type == b"ftyp": if payload_end - payload_start < 4: return None diff --git a/invokeai/app/util/video_audio.py b/invokeai/app/util/video_audio.py index f7cc6eeb1a9..cd7ed1d3fbf 100644 --- a/invokeai/app/util/video_audio.py +++ b/invokeai/app/util/video_audio.py @@ -23,6 +23,9 @@ import numpy as np +# How every soundtrack the app writes is encoded; ``aac_splice`` matches a spliced tail to it. +AAC_ENCODE_ARGS = ["-c:a", "aac"] + def _ffmpeg_exe() -> str: import imageio_ffmpeg @@ -188,8 +191,7 @@ def _run_mux(video_path: Path, wav_path: Path, output_path: Path) -> subprocess. "1:a:0", "-c:v", "copy", - "-c:a", - "aac", + *AAC_ENCODE_ARGS, str(output_path), ], capture_output=True, diff --git a/invokeai/app/util/video_encoding.py b/invokeai/app/util/video_encoding.py index 60c560fcf39..0ae247fe746 100644 --- a/invokeai/app/util/video_encoding.py +++ b/invokeai/app/util/video_encoding.py @@ -10,6 +10,14 @@ yuv420p requires even dimensions; callers validate that before encoding. +Quality and keyframes: frames are encoded at CRF 18 (imageio's default quality maps to +CRF 25), because derived videos — trims, joins, extensions — are routinely re-encoded +from earlier outputs and the losses compound. An IDR keyframe is forced every +``KEYFRAME_INTERVAL_SECONDS`` so a later join can stream-copy the head of this file up to +a keyframe near its end instead of re-encoding it (see ``h264_splice``). Both settings +are part of the SPS/PPS, so changing either makes files written before and after the +change unspliceable with each other: a join between them falls back to re-encoding. + Audio muxing: pass ``audio_path`` to include an audio track. Two constraints: - imageio-ffmpeg does not pass ``-shortest``, so the container duration is the @@ -21,12 +29,16 @@ browsers. """ +import tempfile import wave from pathlib import Path import imageio.v2 as iio2 import numpy as np +CRF = 18 +KEYFRAME_INTERVAL_SECONDS = 2.0 + def make_mp4_writer(path: Path | str, fps: float, audio_path: Path | str | None = None, audio_codec: str = "aac"): """Returns an imageio FFMPEG writer that preserves frame dimensions exactly. @@ -45,7 +57,52 @@ def make_mp4_writer(path: Path | str, fps: float, audio_path: Path | str | None raise FileNotFoundError(f"Audio file for muxing does not exist: {audio_path}") kwargs["audio_path"] = str(audio_path) kwargs["audio_codec"] = audio_codec - return iio2.get_writer(str(path), format="FFMPEG", mode="I", fps=fps, codec="libx264", macro_block_size=1, **kwargs) + keyframe_interval = max(1, round(fps * KEYFRAME_INTERVAL_SECONDS)) + return iio2.get_writer( + str(path), + format="FFMPEG", + mode="I", + fps=fps, + codec="libx264", + macro_block_size=1, + quality=None, + output_params=["-crf", str(CRF), "-g", str(keyframe_interval)], + **kwargs, + ) + + +_writer_decoder_configs: dict[tuple[int, int, float], bytes] = {} + + +def writer_decoder_config(width: int, height: int, fps: float) -> bytes | None: + """The ``avcC`` bytes ``make_mp4_writer`` writes for this canvas and rate. + + They depend only on the encoder settings, canvas and rate, never on the content, so one + encoded frame tells a splice up front whether a file can share a track with new output + — before a long encode is spent finding out. None if the probe cannot be read; only a + successful probe is remembered. + """ + key = (width, height, fps) + if key not in _writer_decoder_configs: + config = _probe_decoder_config(width, height, fps) + if config is None: + return None + _writer_decoder_configs[key] = config + return _writer_decoder_configs[key] + + +def _probe_decoder_config(width: int, height: int, fps: float) -> bytes | None: + from invokeai.app.util.h264_splice import read_h264_layout + + with tempfile.TemporaryDirectory(prefix="invokeai_encoder_probe_") as work: + path = Path(work) / "probe.mp4" + writer = make_mp4_writer(path, fps) + try: + writer.append_data(np.zeros((height, width, 3), dtype=np.uint8)) + finally: + writer.close() + layout = read_h264_layout(path) + return layout.decoder_config if layout is not None else None def write_stereo_wav(path: Path | str, samples: np.ndarray, sample_rate: int) -> None: diff --git a/invokeai/app/util/video_thumbnails.py b/invokeai/app/util/video_thumbnails.py index 7cc13b9d580..ea23a32d81c 100644 --- a/invokeai/app/util/video_thumbnails.py +++ b/invokeai/app/util/video_thumbnails.py @@ -24,7 +24,7 @@ import threading import time from pathlib import Path -from typing import Any, Callable, Iterator, NamedTuple, Optional +from typing import Any, Callable, Generator, Iterator, NamedTuple, Optional import numpy as np import psutil @@ -342,7 +342,7 @@ def iter_video_frames( video_path: Path, timeout: float = VIDEO_DECODE_TIMEOUT_SECONDS, is_canceled: Optional[Callable[[], bool]] = None, -) -> Iterator[np.ndarray]: +) -> Generator[np.ndarray, None, None]: acquired: list[threading.BoundedSemaphore] = [] capacity_deadline = time.monotonic() + timeout try: diff --git a/tests/app/invocations/test_video_concat.py b/tests/app/invocations/test_video_concat.py index 018ddbf5952..60dea24d53f 100644 --- a/tests/app/invocations/test_video_concat.py +++ b/tests/app/invocations/test_video_concat.py @@ -261,13 +261,15 @@ def _invoke(size_mismatch: str, sizes: list[tuple[int, int]]) -> MagicMock: def frames(path: Path, **_: object) -> Iterator[np.ndarray]: w, h = by_path[path.name] - return iter([np.full((h, w, 3), 7, dtype=np.uint8) for _ in range(4)]) + yield from [np.full((h, w, 3), 7, dtype=np.uint8) for _ in range(4)] with ( patch("invokeai.app.invocations.video_concat.probe_video", side_effect=probes), patch("invokeai.app.invocations.video_concat.make_mp4_writer", return_value=MagicMock()), patch("invokeai.app.invocations.video_concat.iter_video_frames", side_effect=frames), patch("invokeai.app.invocations.video_concat.extract_audio_pcm", return_value=None), + # The inputs are not real files, so there is no sample table to splice from. + patch("invokeai.app.invocations.video_concat.read_h264_layout", return_value=None), patch("invokeai.app.invocations.video_concat.VideoOutput.build", return_value=MagicMock()), ): invocation.invoke(context) diff --git a/tests/app/invocations/test_video_concat_splice.py b/tests/app/invocations/test_video_concat_splice.py new file mode 100644 index 00000000000..0f83d656c70 --- /dev/null +++ b/tests/app/invocations/test_video_concat_splice.py @@ -0,0 +1,360 @@ +"""Joining onto a clip this app encoded must not re-encode that clip's head. + +Extend Video joins each continuation onto the video it extends, so a video extended N +times used to carry its first segment through N lossy re-encodes, with artifacts +compounding in the earliest footage. The concat node now stream-copies the first clip up +to its last usable keyframe and encodes only the rest. These tests pin that the copied +frames come through bit-exact (also across repeated extensions), that the joined +timeline has the same length and audio as a full re-encode would, and that inputs which +cannot be spliced losslessly still join correctly by re-encoding. + +Real ffmpeg (the imageio-ffmpeg bundled binary) encodes and decodes throughout. +""" + +import shutil +import struct +import subprocess +from pathlib import Path +from unittest.mock import MagicMock, patch + +import imageio.v2 as iio2 +import imageio_ffmpeg +import numpy as np +import pytest + +from invokeai.app.invocations.fields import VideoField +from invokeai.app.invocations.video_concat import VideoConcatInvocation +from invokeai.app.util.aac_splice import PREROLL_FRAMES +from invokeai.app.util.h264_splice import read_h264_layout +from invokeai.app.util.mp4_metadata import find_box, write_mp4_tags +from invokeai.app.util.video_audio import extract_audio_pcm +from invokeai.app.util.video_encoding import make_mp4_writer, write_stereo_wav +from invokeai.app.util.video_thumbnails import probe_video + +FPS = 16.0 +WIDTH, HEIGHT = 128, 96 +RATE = 32000 + + +def _frames(n: int, seed: int) -> list[np.ndarray]: + # Textured, moving content: flat frames encode near-losslessly and would hide a re-encode. + rng = np.random.default_rng(seed) + texture = rng.integers(0, 256, (HEIGHT + n, WIDTH + n, 3), dtype=np.uint8) + return [texture[i : i + HEIGHT, i : i + WIDTH].copy() for i in range(n)] + + +def _tone(freq: float, n_frames: int) -> np.ndarray: + t = np.arange(round(n_frames / FPS * RATE)) / RATE + mono = (0.5 * np.sin(2 * np.pi * freq * t)).astype(np.float32) + return np.stack([mono, mono]) + + +def _encode(path: Path, frames: list[np.ndarray], tone: float | None = None, legacy: bool = False) -> Path: + audio_path = None + if tone is not None: + audio_path = path.with_suffix(".wav") + write_stereo_wav(audio_path, _tone(tone, len(frames)), RATE) + if legacy: + # A file from before the current encoder settings (CRF 25), with keyframes as frequent as + # ours so a splice point exists and only the decoder configuration rules the copy out. + writer = iio2.get_writer( + str(path), + format="FFMPEG", + mode="I", + fps=FPS, + codec="libx264", + macro_block_size=1, + quality=None, + output_params=["-crf", "25", "-g", "32"], + ) + else: + writer = make_mp4_writer(path, FPS, audio_path=audio_path) + for frame in frames: + writer.append_data(frame) + writer.close() + return path + + +def _decode(path: Path) -> list[np.ndarray]: + reader = iio2.get_reader(str(path), format="FFMPEG") + try: + return [np.asarray(frame) for frame in reader] + finally: + reader.close() + + +def _movie_duration(path: Path) -> float: + """The container's own duration (``mvhd``), at its full precision.""" + with open(path, "rb") as fh: + fh.seek(0, 2) + moov = find_box(fh, 0, fh.tell(), b"moov") + assert moov is not None + mvhd = find_box(fh, *moov, b"mvhd") + assert mvhd is not None + fh.seek(mvhd[0]) + wide = fh.read(1)[0] == 1 + fh.seek(mvhd[0] + (20 if wide else 12)) + timescale, duration = struct.unpack(">IQ" if wide else ">II", fh.read(12 if wide else 8)) + return duration / timescale + + +def _dominant_freq(pcm: np.ndarray, rate: int) -> float: + spectrum = np.abs(np.fft.rfft(pcm[0])) + spectrum[0] = 0.0 + return float(np.fft.rfftfreq(pcm.shape[1], 1.0 / rate)[int(np.argmax(spectrum))]) + + +def _concat(tmp_path: Path, inputs: list[Path], out_name: str, **fields: object) -> tuple[Path, MagicMock]: + """Run the node on real files; returns the saved MP4 and the mocked context.""" + by_name = {p.stem: p for p in inputs} + saved = tmp_path / f"{out_name}.mp4" + + def _save(source_path: Path, **_: object) -> MagicMock: + # The stream-copy remux save() runs to embed the record, so chained joins read what + # production stores. + write_mp4_tags(source_path, saved, {"invokeai_metadata": "{}"}) + return MagicMock() + + context = MagicMock() + context.videos.get_path.side_effect = lambda name: by_name[name] + context.util.is_canceled.return_value = False + context.videos.save.side_effect = _save + node = VideoConcatInvocation(videos=[VideoField(video_name=p.stem) for p in inputs], **fields) # type: ignore[arg-type] + with patch("invokeai.app.invocations.video_concat.VideoOutput.build", return_value=MagicMock()): + node.invoke(context) + return saved, context + + +def _logged(context: MagicMock) -> str: + return "\n".join(str(call.args[0]) for call in context.logger.info.call_args_list) + + +def _last_usable_keyframe(path: Path, tail: int) -> int: + layout = read_h264_layout(path) + assert layout is not None + return max(k for k in layout.sync_samples if 0 < k <= layout.sample_count - max(tail, 1)) + + +@pytest.mark.parametrize( + "transition,transition_frames,expected_len", + [("cut", 0, 100 + 40), ("crossfade", 2, 100 + 40 - 2), ("fade_through_black", 3, 100 + 40)], +) +def test_first_clip_head_is_copied_bit_exact(tmp_path, transition, transition_frames, expected_len): + a = _encode(tmp_path / "a.mp4", _frames(100, seed=1)) + b = _encode(tmp_path / "b.mp4", _frames(40, seed=2)) + out, context = _concat(tmp_path, [a, b], "out", transition=transition, transition_frames=transition_frames) + + tail = {"crossfade": transition_frames, "fade_through_black": transition_frames // 2}.get(transition, 0) + keyframe = _last_usable_keyframe(a, tail) + a_frames, out_frames = _decode(a), _decode(out) + assert len(out_frames) == expected_len + assert all(np.array_equal(x, y) for x, y in zip(a_frames[:keyframe], out_frames[:keyframe], strict=True)) + assert f"Stream-copied the first {keyframe} frames" in _logged(context) + saved = context.videos.save.call_args.kwargs + assert saved["duration"] == pytest.approx(expected_len / FPS) + + +def test_repeated_extension_leaves_the_original_segment_untouched(tmp_path): + a = _encode(tmp_path / "a.mp4", _frames(100, seed=1)) + original = _decode(a) + current = a + for round_ in range(3): + nxt = _encode(tmp_path / f"n{round_}.mp4", _frames(40, seed=10 + round_)) + current, _ = _concat(tmp_path, [current, nxt], f"round{round_}", transition="crossfade", transition_frames=2) + final = _decode(current) + assert len(final) == 100 + 3 * (40 - 2) + # Each round re-encodes only from the last keyframe of its first clip, which lies in the + # newest continuation from the second round on: A's own head is copied every time. + keyframe = _last_usable_keyframe(a, 2) + assert all(np.array_equal(x, y) for x, y in zip(original[:keyframe], final[:keyframe], strict=True)) + + +def test_spliced_join_keeps_the_audio_timeline(tmp_path): + a = _encode(tmp_path / "a.mp4", _frames(64, seed=1), tone=440.0) + b = _encode(tmp_path / "b.mp4", _frames(64, seed=2), tone=880.0) + out, context = _concat(tmp_path, [a, b], "out", transition="cut", transition_frames=0) + assert "Stream-copied" in _logged(context) + + extracted = extract_audio_pcm(out) + assert extracted is not None + pcm, rate = extracted + assert pcm.shape[1] / rate == pytest.approx(128 / FPS, abs=0.1) + half, margin = pcm.shape[1] // 2, rate // 10 + assert _dominant_freq(pcm[:, : half - margin], rate) == pytest.approx(440.0, abs=20.0) + assert _dominant_freq(pcm[:, half + margin :], rate) == pytest.approx(880.0, abs=20.0) + + +def test_first_clip_from_another_encoder_is_re_encoded(tmp_path): + a = _encode(tmp_path / "a.mp4", _frames(100, seed=1), legacy=True) + b = _encode(tmp_path / "b.mp4", _frames(40, seed=2)) + out, context = _concat(tmp_path, [a, b], "out", transition="crossfade", transition_frames=2) + assert "Stream-copied" not in _logged(context) + assert len(_decode(out)) == 138 + + +def test_retimed_join_is_re_encoded(tmp_path): + # A stream copy cannot change the frame rate, so an fps override never splices. + a = _encode(tmp_path / "a.mp4", _frames(100, seed=1)) + b = _encode(tmp_path / "b.mp4", _frames(40, seed=2)) + out, context = _concat(tmp_path, [a, b], "out", transition="cut", fps=24) + assert "Stream-copied" not in _logged(context) + assert len(_decode(out)) == 140 + + +def _remux(src: Path, dst: Path, *args: str) -> Path: + subprocess.run( + [imageio_ffmpeg.get_ffmpeg_exe(), "-y", "-loglevel", "error", *args, "-i", str(src), "-c", "copy", str(dst)], + check=True, + capture_output=True, + ) + return dst + + +def test_stream_copied_trim_as_first_clip_joins(tmp_path): + """A lossless trim (LosslessCut, `ffmpeg -c copy`) keeps every sample and hides the + leading ones with an edit list, so its sample indices are not its frame indices.""" + source = _encode(tmp_path / "source.mp4", _frames(68, seed=1)) + a = _remux(source, tmp_path / "a.mp4", "-ss", "0.3") + b = _encode(tmp_path / "b.mp4", _frames(40, seed=2)) + a_frames = _decode(a) + out, context = _concat(tmp_path, [a, b], "out", transition="cut", transition_frames=0) + assert "Stream-copied" not in _logged(context) + assert len(_decode(out)) == len(a_frames) + 40 + + +def test_rotated_square_first_clip_keeps_its_orientation(tmp_path): + """Rotation metadata leaves a square clip's encoding unchanged, but decoding applies it: + a copied head would carry the rotation onto every re-encoded (already upright) frame.""" + square = _frames(70, seed=1) + square = [f[:, : f.shape[0]].copy() for f in square] + source = tmp_path / "sq.mp4" + writer = make_mp4_writer(source, FPS) + for frame in square: + writer.append_data(frame) + writer.close() + a = _remux(source, tmp_path / "a.mp4", "-display_rotation", "90") + b_frames = [f[:, : f.shape[0]].copy() for f in _frames(20, seed=2)] + b = tmp_path / "b.mp4" + writer = make_mp4_writer(b, FPS) + for frame in b_frames: + writer.append_data(frame) + writer.close() + + out, context = _concat(tmp_path, [a, b], "out", transition="cut", transition_frames=0) + assert "Stream-copied" not in _logged(context) + out_frames, a_frames = _decode(out), _decode(a) + assert len(out_frames) == 90 + # Every output frame matches what its input decodes to (A rotated upright, B as-is). + for got, want in [(out_frames[0], a_frames[0]), (out_frames[-1], _decode(b)[-1])]: + assert np.abs(got.astype(int) - want.astype(int)).mean() < 8 + + +def test_short_remainder_keeps_one_timeline(tmp_path): + """A remainder of 1-2 frames would be encoded with a shorter B-frame delay than the + copied head, distorting the joined timestamps.""" + b = _encode(tmp_path / "b.mp4", _frames(1, seed=2)) + # Every first-clip length just past a keyframe, so one leaves a 1-2 frame remainder. + for n_a in range(65, 70): + a = _encode(tmp_path / f"a{n_a}.mp4", _frames(n_a, seed=1)) + out, _ = _concat(tmp_path, [a, b], f"out{n_a}", transition="cut", transition_frames=0) + _, _, duration, fps = probe_video(out) + assert fps == pytest.approx(FPS, rel=1e-3), f"first clip of {n_a} frames" + assert duration == pytest.approx((n_a + 1) / FPS, abs=0.5 / FPS) + assert len(_decode(out)) == n_a + 1 + + +def test_trimmed_source_extends(tmp_path): + """The Extend graph: trim the source, then join a continuation onto the trim.""" + from invokeai.app.invocations.video_frame_extract_range import ExtractVideoRangeInvocation + + source = _encode(tmp_path / "source.mp4", _frames(100, seed=1)) + trimmed = tmp_path / "trimmed.mp4" + + def _save(source_path: Path, **_: object) -> MagicMock: + shutil.copyfile(source_path, trimmed) + return MagicMock() + + context = MagicMock() + context.videos.get_path.return_value = source + context.util.is_canceled.return_value = False + context.videos.save.side_effect = _save + node = ExtractVideoRangeInvocation(video=VideoField(video_name="source"), start_frame=31, end_frame=66) + base = MagicMock(video=VideoField(video_name="trimmed"), width=WIDTH, height=HEIGHT, num_frames=36, fps=FPS) + with patch("invokeai.app.invocations.video_frame_extract_range.VideoOutput.build", return_value=base): + node.invoke(context) + + b = _encode(tmp_path / "b.mp4", _frames(40, seed=2)) + out, _ = _concat(tmp_path, [trimmed, b], "out", transition="crossfade", transition_frames=2) + assert len(_decode(out)) == 36 + 40 - 2 + + +def test_repeated_extension_leaves_the_original_soundtrack_untouched(tmp_path): + """Every join used to decode and re-encode the whole soundtrack, compounding AAC loss + in the earliest audio just as in its video.""" + a = _encode(tmp_path / "a.mp4", _frames(100, seed=1), tone=440.0) + original = extract_audio_pcm(a) + assert original is not None + current = a + for round_ in range(3): + nxt = _encode(tmp_path / f"n{round_}.mp4", _frames(40, seed=10 + round_), tone=660.0 + 110 * round_) + current, context = _concat(tmp_path, [current, nxt], f"r{round_}", transition="crossfade", transition_frames=2) + assert "audio samples of the first clip without re-encoding" in _logged(context) + frames = 100 + (round_ + 1) * (40 - 2) + # The soundtrack ends with the video, not a padded AAC frame later (a frozen last frame). + assert _movie_duration(current) == pytest.approx(frames / FPS, abs=0.002) + # Decoding yields whole AAC frames: the timeline's samples, padded to less than one more. + decoded = extract_audio_pcm(current) + assert decoded is not None + assert 0 <= decoded[0].shape[1] - round(frames * RATE / FPS) < 1024 + + final = extract_audio_pcm(current) + assert final is not None and final[1] == original[1] == RATE + # A's audio is copied every round up to a whole AAC frame at least PREROLL_FRAMES clear of + # its outgoing crossfade. + unchanged = (100 - 2) * RATE // int(FPS) - (PREROLL_FRAMES + 1) * 1024 + assert np.array_equal(final[0][:, :unchanged], original[0][:, :unchanged]) + # The new audio is where it belongs: the last round's tone closes the soundtrack. + assert _dominant_freq(final[0][:, -RATE // 2 :], RATE) == pytest.approx(880.0, abs=20.0) + + +def test_soundtrack_from_another_encoder_is_re_encoded(tmp_path): + a_video = _encode(tmp_path / "av.mp4", _frames(64, seed=1)) + a_wav = tmp_path / "a.wav" + write_stereo_wav(a_wav, _tone(440.0, 64), RATE) + a = tmp_path / "a.mp4" + # Mono AAC: a different AudioSpecificConfig from the stereo tracks this app writes. + subprocess.run( + [imageio_ffmpeg.get_ffmpeg_exe(), "-y", "-loglevel", "error", "-i", str(a_video), "-i", str(a_wav)] + + ["-map", "0:v", "-map", "1:a", "-c:v", "copy", "-c:a", "aac", "-ac", "1", str(a)], + check=True, + ) + b = _encode(tmp_path / "b.mp4", _frames(64, seed=2), tone=880.0) + out, context = _concat(tmp_path, [a, b], "out", transition="cut", transition_frames=0) + assert "audio samples of the first clip" not in _logged(context) + extracted = extract_audio_pcm(out) + assert extracted is not None + assert extracted[0].shape[1] / RATE == pytest.approx(128 / FPS, abs=0.1) + + +@pytest.mark.parametrize("b_tone", [660.0, None], ids=["b-audible", "b-silent"]) +@pytest.mark.parametrize( + "transition,transition_frames", [("cut", 0), ("crossfade", 2), ("crossfade", 8), ("fade_through_black", 3)] +) +def test_spliced_soundtrack_matches_a_full_re_encode(tmp_path, monkeypatch, transition, transition_frames, b_tone): + """Copying the head must neither move what follows nor glitch the seam: the spliced + soundtrack matches a full re-encode of the same timeline everywhere, seam included, to + within the coding noise of two lossy encodes. A silent continuation (a Wan clip) must + splice too: its silence codes with no error and is no measure of a seam.""" + a = _encode(tmp_path / "a.mp4", _frames(100, seed=1), tone=440.0) + b = _encode(tmp_path / "b.mp4", _frames(40, seed=2), tone=b_tone) + fields = {"transition": transition, "transition_frames": transition_frames} + spliced, context = _concat(tmp_path, [a, b], "spliced", **fields) + assert "audio samples of the first clip without re-encoding" in _logged(context) + + monkeypatch.setattr(VideoConcatInvocation, "_mux_spliced_audio", lambda *args, **kwargs: False) + full, _ = _concat(tmp_path, [a, b], "full", **fields) + + got, want = extract_audio_pcm(spliced), extract_audio_pcm(full) + assert got is not None and want is not None + n = min(got[0].shape[1], want[0].shape[1]) + assert float(np.abs(got[0][:, :n] - want[0][:, :n]).max()) < 0.1 diff --git a/tests/app/invocations/test_video_frame_extract_range.py b/tests/app/invocations/test_video_frame_extract_range.py index c773312c673..debbaa4f68f 100644 --- a/tests/app/invocations/test_video_frame_extract_range.py +++ b/tests/app/invocations/test_video_frame_extract_range.py @@ -7,16 +7,22 @@ stops decoding as soon as the range is written. """ +import shutil from pathlib import Path from typing import Iterator from unittest.mock import MagicMock, patch +import imageio.v2 as iio2 import numpy as np import pytest from invokeai.app.invocations.fields import VideoField from invokeai.app.invocations.video_frame_extract_range import ExtractVideoRangeInvocation, _write_frame_range from invokeai.app.services.session_processor.session_processor_common import CanceledException +from invokeai.app.util.h264_splice import read_h264_layout +from invokeai.app.util.mp4_metadata import read_mp4_tags, write_mp4_tags +from invokeai.app.util.video_encoding import make_mp4_writer +from invokeai.app.util.video_thumbnails import probe_video class RecordingWriter: @@ -98,10 +104,11 @@ def test_invocation_only_saves_complete_requested_range(written: int, should_rai with ( patch("invokeai.app.invocations.video_frame_extract_range.probe_video", return_value=(32, 32, 0.625, 8.0)), - patch("invokeai.app.invocations.video_frame_extract_range.decoder_frame_count", return_value=5), + patch("invokeai.app.invocations.video_frame_extract_range.decoder_frame_count", return_value=10), patch("invokeai.app.invocations.video_frame_extract_range.make_mp4_writer", return_value=MagicMock()), patch("invokeai.app.invocations.video_frame_extract_range._write_frame_range", return_value=written), patch("invokeai.app.invocations.video_frame_extract_range.extract_audio_pcm", return_value=None), + patch("invokeai.app.invocations.video_frame_extract_range.read_h264_layout", return_value=None), patch("invokeai.app.invocations.video_frame_extract_range.VideoOutput.build", return_value=base_output), ): if should_raise: @@ -130,13 +137,208 @@ def test_invocation_echoes_the_source_video_it_was_given() -> None: with ( patch("invokeai.app.invocations.video_frame_extract_range.probe_video", return_value=(32, 32, 0.625, 8.0)), - patch("invokeai.app.invocations.video_frame_extract_range.decoder_frame_count", return_value=5), + patch("invokeai.app.invocations.video_frame_extract_range.decoder_frame_count", return_value=10), patch("invokeai.app.invocations.video_frame_extract_range.make_mp4_writer", return_value=MagicMock()), patch("invokeai.app.invocations.video_frame_extract_range._write_frame_range", return_value=5), patch("invokeai.app.invocations.video_frame_extract_range.extract_audio_pcm", return_value=None), + patch("invokeai.app.invocations.video_frame_extract_range.read_h264_layout", return_value=None), patch("invokeai.app.invocations.video_frame_extract_range.VideoOutput.build", return_value=base_output), ): output = invocation.invoke(context) assert output.source_video.video_name == "input.mp4" assert output.video.video_name == "output.mp4" + + +def _decode(path: Path) -> list[np.ndarray]: + reader = iio2.get_reader(str(path), format="FFMPEG") + try: + return [np.asarray(frame) for frame in reader] + finally: + reader.close() + + +@pytest.mark.parametrize("fps", [0, 16]) +def test_whole_clip_at_its_own_rate_is_copied_without_re_encoding(tmp_path: Path, fps: int) -> None: + """Extend Video passes an untrimmed source through this node on every extension; a + re-encode there would cost the whole source one more lossy generation each time.""" + rng = np.random.default_rng(0) + source = tmp_path / "source.mp4" + writer = make_mp4_writer(source, 16.0) + for _ in range(24): + writer.append_data(rng.integers(0, 256, (48, 64, 3), dtype=np.uint8)) + writer.close() + + saved = tmp_path / "saved.mp4" + context = MagicMock() + context.videos.get_path.return_value = source + context.util.is_canceled.return_value = False + + def _save(source_path: Path, **_: object) -> MagicMock: + shutil.copyfile(source_path, saved) + return MagicMock() + + context.videos.save.side_effect = _save + invocation = ExtractVideoRangeInvocation( + video=VideoField(video_name="source.mp4"), start_frame=0, end_frame=-1, fps=fps + ) + base_output = MagicMock( + video=VideoField(video_name="saved.mp4"), width=64, height=48, num_frames=24, fps=16.0, duration=1.5 + ) + with patch("invokeai.app.invocations.video_frame_extract_range.VideoOutput.build", return_value=base_output): + output = invocation.invoke(context) + + source_frames, saved_frames = _decode(source), _decode(saved) + assert len(saved_frames) == 24 + assert all(np.array_equal(a, b) for a, b in zip(source_frames, saved_frames, strict=True)) + assert (output.start_frame, output.end_frame) == (0, 23) + assert context.videos.save.call_args.kwargs["duration"] == pytest.approx(24 / 16.0) + + +def _textured_source(path: Path, n: int) -> list[np.ndarray]: + rng = np.random.default_rng(1) + texture = rng.integers(0, 256, (48 + n, 64 + n, 3), dtype=np.uint8) + frames = [texture[i : i + 48, i : i + 64].copy() for i in range(n)] + writer = make_mp4_writer(path, 16.0) + for frame in frames: + writer.append_data(frame) + writer.close() + return frames + + +def _invoke_trim(tmp_path: Path, source: Path, start: int, end: int) -> tuple[Path, MagicMock]: + saved = tmp_path / "saved.mp4" + + def _save(source_path: Path, **_: object) -> MagicMock: + shutil.copyfile(source_path, saved) + return MagicMock() + + context = MagicMock() + context.videos.get_path.return_value = source + context.util.is_canceled.return_value = False + context.videos.save.side_effect = _save + invocation = ExtractVideoRangeInvocation( + video=VideoField(video_name="source.mp4"), start_frame=start, end_frame=end + ) + base_output = MagicMock( + video=VideoField(video_name="saved.mp4"), width=64, height=48, num_frames=1, fps=16.0, duration=1.0 + ) + with patch("invokeai.app.invocations.video_frame_extract_range.VideoOutput.build", return_value=base_output): + invocation.invoke(context) + return saved, context + + +@pytest.mark.parametrize("start,end", [(0, 80), (10, 80), (10, 99), (32, 63)]) +def test_trim_copies_the_keyframe_aligned_frames_inside_the_range(tmp_path: Path, start: int, end: int) -> None: + """Trimming the source before extending it must not re-encode the frames that survive + the trim, apart from the partial keyframe intervals at its ends.""" + source = tmp_path / "source.mp4" + _textured_source(source, 100) + layout = read_h264_layout(source) + assert layout is not None + k1 = min(k for k in layout.sync_samples if k >= start) + k2 = max(k for k in [*layout.sync_samples, layout.sample_count] if k <= end + 1) + assert k1 < k2, "fixture must contain a keyframe interval inside the range" + + saved, context = _invoke_trim(tmp_path, source, start, end) + + source_frames, saved_frames = _decode(source), _decode(saved) + assert len(saved_frames) == end - start + 1 + copied = slice(k1 - start, k2 - start) + assert all(np.array_equal(a, b) for a, b in zip(source_frames[k1:k2], saved_frames[copied], strict=True)) + # The re-encoded ends are the right frames, just not bit-exact. + for offset in (0, len(saved_frames) - 1): + error = np.abs(saved_frames[offset].astype(int) - source_frames[start + offset].astype(int)).mean() + assert error < 8 + assert f"Stream-copied frames {k1}-{k2 - 1}" in str(context.logger.info.call_args_list) + assert context.videos.save.call_args.kwargs["duration"] == pytest.approx((end - start + 1) / 16.0) + + +def test_trim_within_one_keyframe_interval_is_re_encoded(tmp_path: Path) -> None: + source = tmp_path / "source.mp4" + _textured_source(source, 100) + saved, context = _invoke_trim(tmp_path, source, 5, 20) + assert len(_decode(saved)) == 16 + assert "Stream-copied" not in str(context.logger.info.call_args_list) + + +def test_whole_clip_copy_does_not_carry_the_sources_generation_metadata(tmp_path: Path) -> None: + """save() rewrites the record into the file, but falls back to a sidecar if that fails, + and embedded tags win on read: a copy still carrying the source's would be read back + as this output's.""" + plain = tmp_path / "plain.mp4" + _textured_source(plain, 24) + source = tmp_path / "source.mp4" + write_mp4_tags(plain, source, {"invokeai_metadata": '{"from": "source"}'}) + assert "invokeai_metadata" in read_mp4_tags(source) + + saved, _ = _invoke_trim(tmp_path, source, 0, -1) + assert "invokeai_metadata" not in read_mp4_tags(saved) + assert len(_decode(saved)) == 24 + + +@pytest.mark.parametrize("start,end", [(31, 66), (30, 63), (32, 64), (32, 65), (31, 64), (63, 65)]) +def test_short_encoded_ends_keep_one_timeline(tmp_path: Path, start: int, end: int) -> None: + """x264 encodes a 1-2 frame part with a shorter B-frame delay than a full one, and joining + it to a copied run squeezed the timestamps: the trim probed at the wrong fps, and the + extension built on it was refused for mismatched frame rates.""" + source = tmp_path / "source.mp4" + _textured_source(source, 100) + saved, _ = _invoke_trim(tmp_path, source, start, end) + + n = end - start + 1 + width, height, duration, fps = probe_video(saved) + assert fps == pytest.approx(16.0, rel=1e-3) + assert duration == pytest.approx(n / 16.0, abs=0.5 / 16.0) + assert len(_decode(saved)) == n + + +def test_open_gop_source_trim_keeps_every_frame(tmp_path: Path) -> None: + """An open-GOP I-frame is a sync sample but not a clean cut: the frames after it in + decode order that display before it would be lost from a copied run.""" + rng = np.random.default_rng(3) + texture = rng.integers(0, 256, (48 + 100, 64 + 100, 3), dtype=np.uint8) + frames = [texture[i : i + 48, i : i + 64].copy() for i in range(100)] + source = tmp_path / "source.mp4" + writer = iio2.get_writer( + str(source), + format="FFMPEG", + mode="I", + fps=16.0, + codec="libx264", + macro_block_size=1, + quality=None, + output_params=["-crf", "18", "-g", "32", "-x264-params", "open-gop=1"], + ) + for frame in frames: + writer.append_data(frame) + writer.close() + + # Every start in the last two GOPs, so one lands on an open-GOP I-frame wherever x264 put it. + for start in range(60, 100): + saved, _ = _invoke_trim(tmp_path, source, start, 99) + assert len(_decode(saved)) == 100 - start, f"trim from {start}" + + +@pytest.mark.parametrize("start,end", [(31, 130), (40, 129)]) +def test_cut_moves_inward_rather_than_leave_a_short_encoded_end(tmp_path: Path, start: int, end: int) -> None: + source = tmp_path / "source.mp4" + _textured_source(source, 200) + saved, context = _invoke_trim(tmp_path, source, start, end) + assert "Stream-copied frames" in str(context.logger.info.call_args_list) + _, _, _, fps = probe_video(saved) + assert fps == pytest.approx(16.0, rel=1e-3) + assert len(_decode(saved)) == end - start + 1 + + +def test_squeezed_join_falls_back_to_re_encoding(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """The joined file's timeline is checked, not assumed: with the short-end rule disabled, + a 1-frame encoded head squeezes the timestamps and the node must re-encode instead.""" + monkeypatch.setattr("invokeai.app.invocations.video_frame_extract_range.MIN_ENCODED_FRAMES", 1) + source = tmp_path / "source.mp4" + _textured_source(source, 100) + saved, context = _invoke_trim(tmp_path, source, 31, 95) + assert "Re-encoding the whole range" in str(context.logger.info.call_args_list) + _, _, _, fps = probe_video(saved) + assert fps == pytest.approx(16.0, rel=1e-3) + assert len(_decode(saved)) == 65 diff --git a/tests/app/util/test_h264_splice.py b/tests/app/util/test_h264_splice.py new file mode 100644 index 00000000000..8612b155ade --- /dev/null +++ b/tests/app/util/test_h264_splice.py @@ -0,0 +1,103 @@ +"""read_h264_layout must refuse tracks whose sample indices are not the decoder's frames. + +A splice copies samples by index and re-encodes the neighbouring frames from a decode, so +the two must name the same frames. A ``-c copy`` trim keeps every sample and hides some +with an edit list, and a metadata rotation is applied to decoded frames but carried as-is +by a stream copy; either would misalign or misorient a join. Real ffmpeg throughout. +""" + +import subprocess +from pathlib import Path + +import imageio.v2 as iio2 +import imageio_ffmpeg +import numpy as np +import pytest + +from invokeai.app.util.h264_splice import read_h264_layout +from invokeai.app.util.video_encoding import make_mp4_writer, writer_decoder_config + +FPS = 16.0 + + +def _encode(path: Path, n: int, width: int = 96, height: int = 64) -> Path: + rng = np.random.default_rng(0) + writer = make_mp4_writer(path, FPS) + for _ in range(n): + writer.append_data(rng.integers(0, 256, (height, width, 3), dtype=np.uint8)) + writer.close() + return path + + +def _remux(src: Path, dst: Path, *args: str) -> Path: + subprocess.run( + [imageio_ffmpeg.get_ffmpeg_exe(), "-y", "-loglevel", "error", *args, "-i", str(src), "-c", "copy", str(dst)], + check=True, + capture_output=True, + ) + return dst + + +def _decoded_count(path: Path) -> int: + reader = iio2.get_reader(str(path), format="FFMPEG") + try: + return sum(1 for _ in reader) + finally: + reader.close() + + +def test_own_output_lists_exactly_the_decoded_frames(tmp_path: Path) -> None: + path = _encode(tmp_path / "a.mp4", 70) + layout = read_h264_layout(path) + assert layout is not None + assert layout.sample_count == _decoded_count(path) == 70 + assert layout.sync_samples[0] == 0 and len(layout.sync_samples) > 1 + assert layout.decoder_config == writer_decoder_config(96, 64, FPS) + + +def test_edit_list_trim_is_refused(tmp_path: Path) -> None: + # A stream-copied trim starting off a keyframe keeps every sample from the keyframe before + # it and hides the leading ones with an edit list. + trimmed = _remux(_encode(tmp_path / "a.mp4", 70), tmp_path / "trimmed.mp4", "-ss", "0.3") + assert _decoded_count(trimmed) < 70 + assert read_h264_layout(trimmed) is None + + +@pytest.mark.parametrize("trim", [("-ss", "0.3"), ("-ss", "2.0"), ("-t", "2.1"), ("-ss", "0.5", "-t", "1.3")]) +def test_an_accepted_layout_lists_exactly_the_decoded_frames(tmp_path: Path, trim: tuple[str, ...]) -> None: + trimmed = _remux(_encode(tmp_path / "a.mp4", 70), tmp_path / "trimmed.mp4", *trim) + layout = read_h264_layout(trimmed) + assert layout is None or layout.sample_count == _decoded_count(trimmed) + + +def test_rotation_metadata_is_refused(tmp_path: Path) -> None: + source = _encode(tmp_path / "sq.mp4", 40, width=64, height=64) + rotated = _remux(source, tmp_path / "rot.mp4", "-display_rotation", "90") + assert read_h264_layout(rotated) is None + + +def test_non_mp4_is_refused(tmp_path: Path) -> None: + path = tmp_path / "x.mp4" + path.write_bytes(b"not a video") + assert read_h264_layout(path) is None + + +def test_open_gop_i_frames_are_not_cut_points(tmp_path: Path) -> None: + path = tmp_path / "og.mp4" + rng = np.random.default_rng(0) + writer = iio2.get_writer( + str(path), + format="FFMPEG", + mode="I", + fps=FPS, + codec="libx264", + macro_block_size=1, + quality=None, + output_params=["-g", "32", "-x264-params", "open-gop=1"], + ) + for _ in range(100): + writer.append_data(rng.integers(0, 256, (64, 96, 3), dtype=np.uint8)) + writer.close() + layout = read_h264_layout(path) + assert layout is not None + assert layout.sync_samples == (0,) From 9a607ccf05b4cac39832bb559286771107216cce Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Mon, 12 Oct 2026 23:34:02 -0400 Subject: [PATCH 2/7] fix(video): bound splice parsing, compare whole sample entries, cancel splices Sample-table and sample-entry reads are bounded before allocating; eligibility compares the full sample entry (colr/pasp/...) and the joined file must keep it. ffmpeg splice children are killed on cancellation, and nodes re-check cancellation before saving. --- invokeai/app/invocations/video_concat.py | 26 +- .../invocations/video_frame_extract_range.py | 21 +- invokeai/app/util/aac_splice.py | 21 +- invokeai/app/util/h264_splice.py | 229 +++++++++++++----- invokeai/app/util/video_encoding.py | 19 +- .../invocations/test_video_concat_splice.py | 73 ++++++ .../test_video_frame_extract_range.py | 28 +++ tests/app/util/test_h264_splice.py | 162 ++++++++++++- 8 files changed, 494 insertions(+), 85 deletions(-) diff --git a/invokeai/app/invocations/video_concat.py b/invokeai/app/invocations/video_concat.py index 8ddeff8ba0b..0701a488b09 100644 --- a/invokeai/app/invocations/video_concat.py +++ b/invokeai/app/invocations/video_concat.py @@ -63,6 +63,7 @@ SpliceError, concat_stream_copy, copy_frames, + probes_as, read_h264_layout, ) from invokeai.app.util.video_audio import ( @@ -71,7 +72,7 @@ mux_audio_into_video, resample_linear, ) -from invokeai.app.util.video_encoding import make_mp4_writer, write_stereo_wav, writer_decoder_config +from invokeai.app.util.video_encoding import make_mp4_writer, write_stereo_wav, writer_sample_entry from invokeai.app.util.video_thumbnails import iter_video_frames, probe_video TransitionMode = Literal["cut", "crossfade", "fade_through_black"] @@ -235,7 +236,9 @@ def invoke(self, context: InvocationContext) -> VideoOutput: keyframe, layout = splice rest_path = work_dir / "rest.mp4" rest = self._encode_joined(context, paths, (width, height), output_fps, rest_path, skip=keyframe) - if self._splice(context, paths[0], keyframe, layout, rest_path, rest[0][0], work_dir, tmp_path): + if self._splice( + context, paths[0], keyframe, layout, rest_path, rest[0][0], output_fps, work_dir, tmp_path + ): joined = rest if joined is None: joined = self._encode_joined(context, paths, (width, height), output_fps, tmp_path, skip=0) @@ -287,6 +290,9 @@ def _save_joined( f"Encoded concatenated MP4: {num_frames} frames @ {output_fps:.2f} fps " f"({duration:.2f}s) at {width}x{height}" + (", with AAC audio" if audio is not None else ", silent") ) + # Splice stages run uninterrupted; a cancellation that arrived meanwhile must not be saved. + if context.util.is_canceled(): + raise CanceledException video_dto = context.videos.save( source_path=source_path, width=width, @@ -324,7 +330,7 @@ def _plan_splice( if ( layout is None or layout.frame_duration is None - or layout.decoder_config != writer_decoder_config(*canvas, output_fps) + or layout.sample_entry != writer_sample_entry(*canvas, output_fps) ): return None tail_need, _ = self._transition_windows() @@ -383,6 +389,7 @@ def _splice( layout: H264Layout, rest_path: Path, first_clip_frames: int, + output_fps: float, work_dir: Path, out_path: Path, ) -> bool: @@ -397,24 +404,28 @@ def _splice( rest = read_h264_layout(rest_path) if first_clip_frames != layout.sample_count: reason = f"it decoded to {first_clip_frames} frames but lists {layout.sample_count}" - elif rest is None or rest.decoder_config != layout.decoder_config or rest.sync_samples[:1] != (0,): + elif rest is None or rest.sample_entry != layout.sample_entry or rest.sync_samples[:1] != (0,): reason = "its encoding differs from this node's" else: reason = None if reason is None: head_path = work_dir / "head.mp4" try: - copy_frames(first_path, 0, keyframe, layout.sample_count, head_path) + copy_frames(first_path, 0, keyframe, layout.sample_count, head_path, context.util.is_canceled) head = read_h264_layout(head_path) - if head is None or head.sample_count != keyframe: + # ffmpeg can derive presentation boxes (e.g. a pasp from a tkhd display size) while + # copying, so the copied head must keep the source's own sample entry. + if head is None or head.sample_count != keyframe or head.sample_entry != layout.sample_entry: reason = "copying its head did not yield the expected frames" else: - concat_stream_copy([head_path, rest_path], out_path) + concat_stream_copy([head_path, rest_path], out_path, context.util.is_canceled) joined = read_h264_layout(out_path) if ( joined is None + or joined.sample_entry != layout.sample_entry or joined.sample_count != keyframe + rest.sample_count or joined.frame_duration != layout.frame_duration + or not probes_as(out_path, output_fps, joined.sample_count) ): reason = "the joined file does not hold the expected frames on one timeline" except SpliceError as e: @@ -480,6 +491,7 @@ def _mux_spliced_audio( rate, pcm.shape[1], out_path, + context.util.is_canceled, ) muxed = extract_audio_pcm(out_path) except (SpliceError, AudioExtractionError) as e: diff --git a/invokeai/app/invocations/video_frame_extract_range.py b/invokeai/app/invocations/video_frame_extract_range.py index eb67c19ddd3..fd798478c8b 100644 --- a/invokeai/app/invocations/video_frame_extract_range.py +++ b/invokeai/app/invocations/video_frame_extract_range.py @@ -54,11 +54,12 @@ SpliceError, concat_stream_copy, copy_frames, + probes_as, read_h264_layout, ) from invokeai.app.util.mp4_metadata import write_mp4_tags from invokeai.app.util.video_audio import extract_audio_pcm, mux_audio_into_video, resample_linear -from invokeai.app.util.video_encoding import make_mp4_writer, write_stereo_wav, writer_decoder_config +from invokeai.app.util.video_encoding import make_mp4_writer, write_stereo_wav, writer_sample_entry from invokeai.app.util.video_thumbnails import decoder_frame_count, iter_video_frames, probe_video @@ -295,6 +296,9 @@ def invoke(self, context: InvocationContext) -> ExtractVideoRangeOutput: f"Encoded trimmed MP4: {num_frames} frames @ {output_fps:.2f} fps " f"({out_duration:.2f}s) at {width}x{height}" + (", with AAC audio" if audio is not None else ", silent") ) + # A cancellation that arrived during an uninterruptible stage must not be saved. + if context.util.is_canceled(): + raise CanceledException video_dto = context.videos.save( source_path=source_path, width=width, @@ -338,7 +342,7 @@ def _plan_copy_span( node's encoder writes and a constant frame duration, or the re-encoded ends could not share its track. """ - if layout.frame_duration is None or layout.decoder_config != writer_decoder_config(width, height, fps): + if layout.frame_duration is None or layout.sample_entry != writer_sample_entry(width, height, fps): return None boundaries = [*layout.sync_samples, layout.sample_count] inner_starts = [k for k in layout.sync_samples if k >= start and not 0 < k - start < MIN_ENCODED_FRAMES] @@ -392,25 +396,27 @@ def _write_spliced_range( middle_path = work_dir / "middle.mp4" reason: Optional[str] = None if any( - (part := read_h264_layout(path)) is None or part.decoder_config != layout.decoder_config + (part := read_h264_layout(path)) is None or part.sample_entry != layout.sample_entry for _, _, path in encoded ): reason = "its encoding differs from this node's" else: try: - copy_frames(video_path, k1, k2, layout.sample_count, middle_path) + copy_frames(video_path, k1, k2, layout.sample_count, middle_path, context.util.is_canceled) middle = read_h264_layout(middle_path) - if middle is None or middle.sample_count != k2 - k1: + if middle is None or middle.sample_count != k2 - k1 or middle.sample_entry != layout.sample_entry: reason = "copying its frames did not yield the expected count" else: parts = [head[2]] if head in encoded else [] parts += [middle_path] + ([tail[2]] if tail in encoded else []) - concat_stream_copy(parts, out_path) + concat_stream_copy(parts, out_path, context.util.is_canceled) joined = read_h264_layout(out_path) if ( joined is None + or joined.sample_entry != layout.sample_entry or joined.sample_count != written + (k2 - k1) or joined.frame_duration != layout.frame_duration + or not probes_as(out_path, output_fps, joined.sample_count) ): reason = "the joined file does not hold the expected frames on one timeline" except SpliceError as e: @@ -436,6 +442,9 @@ def _save_unchanged( # A stream-copy remux that drops the source's keyed metadata: save() writes this # node's record, and if that rewrite fails the source's must not be read back as it. write_mp4_tags(video_path, copy_path, {}) + # A cancellation that arrived during an uninterruptible stage must not be saved. + if context.util.is_canceled(): + raise CanceledException video_dto = context.videos.save( source_path=copy_path, width=width, diff --git a/invokeai/app/util/aac_splice.py b/invokeai/app/util/aac_splice.py index e1458635804..ddb8b412219 100644 --- a/invokeai/app/util/aac_splice.py +++ b/invokeai/app/util/aac_splice.py @@ -28,7 +28,7 @@ import struct import tempfile from pathlib import Path -from typing import BinaryIO, NamedTuple, Optional +from typing import BinaryIO, Callable, NamedTuple, Optional import numpy as np @@ -151,7 +151,17 @@ def _read_audio_track(fh: BinaryIO, trak: tuple[int, int], mdia: tuple[int, int] fh.seek(entry_start + 8) if fh.read(2) != b"\x00\x00": # QuickTime v1/v2 sound entries lay their fields out differently return None - esds = find_box(fh, entry_start + _AUDIO_SAMPLE_ENTRY_HEADER, entry_end, b"esds") + # Anything beyond the decoder config (and a per-file bitrate) — a channel layout, say — + # would describe the copied packets differently from the ones encoded beside them. Read + # lazily and bounded: the entry is the uploader's. + # Each allowed box at most once, so a third box of any kind is refused and no more are read. + esds, seen = None, set() + for child, start, end in iter_boxes(fh, entry_start + _AUDIO_SAMPLE_ENTRY_HEADER, entry_end, max_boxes=3): + if child not in (b"esds", b"btrt") or child in seen: + return None + seen.add(child) + if child == b"esds": + esds = (start, end) if esds is None: return None fh.seek(esds[0]) @@ -295,6 +305,7 @@ def mux_spliced_aac( sample_rate: int, total_samples: int, out_path: Path, + is_canceled: Optional[Callable[[], bool]] = None, ) -> np.ndarray: """Write ``out_path`` = ``video_path``'s video plus a soundtrack spliced from ``source``. @@ -316,9 +327,12 @@ def mux_spliced_aac( run_ffmpeg( ["-i", str(source), "-map", "0:a:0", "-c", "copy", "-frames:a", str(head_packets), "-f", "adts", str(head)], f"copying the first {head_packets} audio packets of {source.name}", + is_canceled, ) write_stereo_wav(tail_wav, tail_pcm, sample_rate) - run_ffmpeg(["-i", str(tail_wav), *AAC_ENCODE_ARGS, "-f", "adts", str(tail)], "encoding the new audio") + run_ffmpeg( + ["-i", str(tail_wav), *AAC_ENCODE_ARGS, "-f", "adts", str(tail)], "encoding the new audio", is_canceled + ) head_frames = _adts_frames(head.read_bytes()) if len(head_frames) != head_packets: raise SpliceError(f"copied {len(head_frames)} audio packets, expected {head_packets}") @@ -354,6 +368,7 @@ def mux_spliced_aac( str(out_path), ], "muxing the spliced audio", + is_canceled, ) _set_track_duration(out_path, b"soun", total_samples, sample_rate) return seam_reference diff --git a/invokeai/app/util/h264_splice.py b/invokeai/app/util/h264_splice.py index 83af5c7459f..33712d45dca 100644 --- a/invokeai/app/util/h264_splice.py +++ b/invokeai/app/util/h264_splice.py @@ -6,10 +6,11 @@ :func:`copy_frames` cuts out a keyframe-aligned run of frames without decoding, and :func:`concat_stream_copy` joins it to freshly encoded neighbours. -That join is only valid when every part carries the *same* decoder configuration: an MP4 -stores one ``avcC`` (profile, level, SPS and PPS) per track, so a part encoded with -different parameters would be decoded against the wrong headers. :func:`read_h264_layout` -exposes the ``avcC`` bytes so callers can require an exact match. Parts written by this +That join is only valid when every part carries the *same* sample entry: an MP4 stores +one per track — the ``avcC`` (profile, level, SPS and PPS) a decoder needs, and boxes such +as ``colr`` and ``pasp`` that say how to present the frames — so a part encoded or tagged +differently would be decoded or displayed against the wrong ones. +:func:`read_h264_layout` exposes those bytes so callers can require an exact match. Parts written by this app's own encoder (``make_mp4_writer``) at the same canvas and frame rate match; anything else is reported as not spliceable rather than guessed at. @@ -32,30 +33,42 @@ a stream copy carries into the joined file for all of them). """ +import math import shutil import struct import subprocess import tempfile +import time from pathlib import Path -from typing import BinaryIO, NamedTuple, Optional +from typing import BinaryIO, Callable, NamedTuple, Optional import imageio_ffmpeg +import numpy as np +from invokeai.app.services.session_processor.session_processor_common import CanceledException from invokeai.app.util.mp4_metadata import find_box, iter_boxes # A stream copy is disk-bound; this only stops a wedged child from holding the node forever. SPLICE_TIMEOUT_SECONDS = 600 +# How often a running ffmpeg child is checked for cancellation. +_CANCEL_POLL_SECONDS = 0.1 # ``moov`` precedes or follows ``mdat`` in a real file; a dozen top-level boxes at most. _MAX_TOP_LEVEL_BOXES = 1024 -# A sync-sample table for even an hour of 120 fps all-intra video is ~430k entries. -_MAX_SYNC_ENTRIES = 1 << 20 -# VisualSampleEntry fields before its child boxes (SampleEntry 8 + visual header 70). +# Keyframes examined per track: at the 2 s interval this app writes, nine hours of video. +_MAX_SYNC_ENTRIES = 1 << 14 +# VisualSampleEntry fields before its child boxes (SampleEntry 8 + visual header 70), and where +# among them the 32-byte compressor name sits. _VISUAL_SAMPLE_ENTRY_HEADER = 78 +_COMPRESSOR_NAME = slice(42, 74) +# A real sample entry holds a handful of boxes (avcC, btrt, colr, pasp...); more is hostile. +_MAX_ENTRY_CHILDREN = 32 +# NAL units examined per keyframe before its first slice (AUD, SEI, SPS, PPS precede it). +_MAX_NALS_BEFORE_SLICE = 64 # Shortest part x264 (bframes=3, as make_mp4_writer leaves it) encodes with its full B-frame # reorder delay; a shorter part joined to full-delay parts distorts the joined timestamps. MIN_ENCODED_FRAMES = 3 -# A sample table for an hour of 120 fps video is ~430k entries. -_MAX_SAMPLES = 1 << 22 +# Samples per track: over four hours at 60 fps. Bounds every table read before it is allocated. +_MAX_SAMPLES = 1 << 20 # NAL unit types: 1-5 are coded slices, 5 an IDR picture's. _NAL_IDR = 5 _NAL_VCL = range(1, 6) @@ -72,6 +85,11 @@ class H264Layout(NamedTuple): """0-based decode-order indices of the IDR pictures — the only clean cut points — ascending.""" decoder_config: bytes """The raw ``avcC`` payload: profile, level, SPS and PPS.""" + sample_entry: bytes + """What a joined track inherits from this one's sample entry: its fixed fields and every + child box — ``avcC``, and any ``colr``, ``pasp``, ``clap`` and the like, which change how + frames are presented without changing ``avcC`` — except ``btrt``, a per-file bitrate. + Two tracks can share a joined track only if these bytes match.""" frame_duration: Optional[int] """Every sample's duration in media ticks, or None if they are not all equal.""" @@ -216,58 +234,82 @@ def _read_stts(fh: BinaryIO, stbl: tuple[int, int]) -> Optional[list[tuple[int, return list(zip(flat[::2], flat[1::2], strict=True)) -def _sample_locations(fh: BinaryIO, stbl: tuple[int, int], sample_count: int) -> Optional[list[tuple[int, int]]]: - """``(file_offset, size)`` of every sample, from the size, chunk and chunk-offset tables.""" - stsz, stsc = find_box(fh, *stbl, b"stsz"), find_box(fh, *stbl, b"stsc") - stco, co64 = find_box(fh, *stbl, b"stco"), find_box(fh, *stbl, b"co64") - offsets_box, offset_format = (stco, "I") if stco is not None else (co64, "Q") - if stsz is None or stsc is None or offsets_box is None or sample_count > _MAX_SAMPLES: - return None - # Each table opens with version/flags and a 4-byte count (stsz: a size and a count). - if any(box[0] + 8 > box[1] for box in (stsz, stsc, offsets_box)): +def _sample_sizes(fh: BinaryIO, stsz: tuple[int, int], sample_count: int, file_size: int) -> Optional[np.ndarray]: + """Every sample's size, or None unless the table is plausible for a file of ``file_size``. + + The count comes from the file, so it is bounded before anything is allocated by it: by + ``_MAX_SAMPLES``, by the table actually holding that many entries, and by the samples + fitting in the file. + """ + if not 0 < sample_count <= _MAX_SAMPLES or stsz[0] + 12 > stsz[1]: return None fh.seek(stsz[0] + 4) (uniform_size,) = struct.unpack(">I", fh.read(4)) if uniform_size: - sizes: tuple[int, ...] = (uniform_size,) * sample_count - else: - if stsz[0] + 12 + 4 * sample_count > stsz[1]: + if uniform_size * sample_count > file_size: return None - fh.seek(stsz[0] + 12) - sizes = struct.unpack(f">{sample_count}I", fh.read(4 * sample_count)) + return np.full(sample_count, uniform_size, dtype=np.int64) + if stsz[0] + 12 + 4 * sample_count > stsz[1]: + return None + fh.seek(stsz[0] + 12) + sizes = np.frombuffer(fh.read(4 * sample_count), dtype=">u4").astype(np.int64) + return sizes if int(sizes.sum()) <= file_size else None + + +def _sample_locations( + fh: BinaryIO, stbl: tuple[int, int], sizes: np.ndarray, wanted: tuple[int, ...], file_size: int +) -> Optional[list[tuple[int, int]]]: + """``(file_offset, size)`` of each sample in ``wanted`` (ascending), from the chunk tables.""" + stsc = find_box(fh, *stbl, b"stsc") + stco, co64 = find_box(fh, *stbl, b"stco"), find_box(fh, *stbl, b"co64") + offsets_box, offset_format = (stco, "I") if stco is not None else (co64, "Q") + # Each table opens with version/flags and a 4-byte count. + if stsc is None or offsets_box is None or any(box[0] + 8 > box[1] for box in (stsc, offsets_box)): + return None + sample_count = len(sizes) fh.seek(stsc[0] + 4) (run_count,) = struct.unpack(">I", fh.read(4)) - if run_count == 0 or stsc[0] + 8 + 12 * run_count > stsc[1]: + if not 0 < run_count <= _MAX_SAMPLES or stsc[0] + 8 + 12 * run_count > stsc[1]: return None flat = struct.unpack(f">{3 * run_count}I", fh.read(12 * run_count)) runs = [(flat[i], flat[i + 1]) for i in range(0, len(flat), 3)] fh.seek(offsets_box[0] + 4) (chunk_count,) = struct.unpack(">I", fh.read(4)) entry = struct.calcsize(offset_format) - if chunk_count > _MAX_SAMPLES or offsets_box[0] + 8 + entry * chunk_count > offsets_box[1]: + if chunk_count > sample_count or offsets_box[0] + 8 + entry * chunk_count > offsets_box[1]: return None chunk_offsets = struct.unpack(f">{chunk_count}{offset_format}", fh.read(entry * chunk_count)) + # Byte position of each sample relative to the start of the sample table's data. + starts = np.concatenate([[0], np.cumsum(sizes)]) locations: list[tuple[int, int]] = [] + pending = iter(wanted) + target = next(pending, None) + first_sample = 0 for run, (first_chunk, per_chunk) in enumerate(runs): last_chunk = runs[run + 1][0] - 1 if run + 1 < len(runs) else chunk_count - if not 1 <= first_chunk <= last_chunk + 1 or last_chunk > chunk_count: + if not 1 <= first_chunk <= last_chunk + 1 or last_chunk > chunk_count or per_chunk > sample_count: return None for chunk in range(first_chunk - 1, last_chunk): - position = chunk_offsets[chunk] - for _ in range(per_chunk): - if len(locations) == sample_count: + end_sample = first_sample + per_chunk + if end_sample > sample_count: + return None + while target is not None and target < end_sample: + offset = chunk_offsets[chunk] + int(starts[target] - starts[first_sample]) + if offset + int(sizes[target]) > file_size: return None - size = sizes[len(locations)] - locations.append((position, size)) - position += size - return locations if len(locations) == sample_count else None + locations.append((offset, int(sizes[target]))) + target = next(pending, None) + first_sample = end_sample + return locations if first_sample == sample_count and target is None else None def _is_idr(fh: BinaryIO, offset: int, size: int, length_size: int) -> bool: """True if the sample's first coded slice is an IDR slice (AUD, SEI, SPS and PPS may precede it).""" position, end = offset, offset + size - while position + length_size + 1 <= end: + for _ in range(_MAX_NALS_BEFORE_SLICE): + if position + length_size + 1 > end: + return False fh.seek(position) header = fh.read(length_size + 1) if len(header) < length_size + 1: @@ -303,17 +345,39 @@ def _read_video_stbl(fh: BinaryIO, stbl_start: int, stbl_end: int) -> Optional[H return None fh.seek(avcc[0]) decoder_config = fh.read(avcc[1] - avcc[0]) - - sizes = _read_full_box_table(fh, *stsz, ">II") - if sizes is None: + fh.seek(entry_start) + header_fields = bytearray(fh.read(_VISUAL_SAMPLE_ENTRY_HEADER)) + # The compressor name records the encoder's version: masked, so files this app wrote before + # an ffmpeg upgrade still match what it writes after one. + header_fields[_COMPRESSOR_NAME] = bytes(_COMPRESSOR_NAME.stop - _COMPRESSOR_NAME.start) + parts = [bytes(header_fields)] + children = list( + iter_boxes(fh, entry_start + _VISUAL_SAMPLE_ENTRY_HEADER, entry_end, max_boxes=_MAX_ENTRY_CHILDREN + 1) + ) + if len(children) > _MAX_ENTRY_CHILDREN: return None - sample_count = sizes[1] - if sample_count == 0: + for child, child_start, child_end in children: + if child != b"btrt": + fh.seek(child_start) + parts.append(struct.pack(">I", child_end - child_start) + child + fh.read(child_end - child_start)) + sample_entry = b"".join(parts) + + header = _read_full_box_table(fh, *stsz, ">II") + if header is None: + return None + fh.seek(0, 2) + file_size = fh.tell() + sizes = _sample_sizes(fh, stsz, header[1], file_size) + if sizes is None: return None + sample_count = len(sizes) stss = find_box(fh, *stbl, b"stss") if stss is None: - # No sync-sample table means every sample is a sync sample. + # No sync-sample table means every sample is a sync sample (an all-intra stream), each + # of which would be examined below: only a short one is worth splicing. + if sample_count > _MAX_SYNC_ENTRIES: + return None sync_samples = tuple(range(sample_count)) else: count = _read_full_box_table(fh, *stss, ">I") @@ -328,10 +392,12 @@ def _read_video_stbl(fh: BinaryIO, stbl_start: int, stbl_end: int) -> Optional[H if len(decoder_config) < 5: return None length_size = (decoder_config[4] & 0x03) + 1 - locations = _sample_locations(fh, stbl, sample_count) + locations = _sample_locations(fh, stbl, sizes, sync_samples, file_size) if locations is None: return None - idr_samples = tuple(k for k in sync_samples if _is_idr(fh, *locations[k], length_size)) + idr_samples = tuple( + k for k, location in zip(sync_samples, locations, strict=True) if _is_idr(fh, *location, length_size) + ) runs = _read_stts(fh, stbl) if runs is None: @@ -342,25 +408,56 @@ def _read_video_stbl(fh: BinaryIO, stbl_start: int, stbl_end: int) -> Optional[H sample_count=sample_count, sync_samples=idr_samples, decoder_config=decoder_config, + sample_entry=sample_entry, frame_duration=frame_duration, ) -def run_ffmpeg(args: list[str], description: str) -> None: - try: - proc = subprocess.run( +def run_ffmpeg(args: list[str], description: str, is_canceled: Optional[Callable[[], bool]] = None) -> None: + """Run the bundled ffmpeg, raising SpliceError if it fails or outlasts the timeout. + + ``is_canceled`` is polled while the child runs; a cancellation kills and reaps it and + raises CanceledException, so a canceled node never waits on a long copy. + """ + if is_canceled is not None and is_canceled(): + raise CanceledException + with tempfile.TemporaryFile() as stderr: + proc = subprocess.Popen( [imageio_ffmpeg.get_ffmpeg_exe(), "-y", "-loglevel", "error", *args], - capture_output=True, - timeout=SPLICE_TIMEOUT_SECONDS, + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=stderr, ) - except subprocess.TimeoutExpired as e: - raise SpliceError(f"ffmpeg timed out {description}") from e - if proc.returncode != 0: - stderr = proc.stderr.decode("utf-8", errors="replace").strip() - raise SpliceError(f"ffmpeg failed {description}: {stderr[-500:]}") - - -def copy_frames(src: Path, start: int, end: int, sample_count: int, dst: Path) -> None: + deadline = time.monotonic() + SPLICE_TIMEOUT_SECONDS + try: + while True: + try: + proc.wait(timeout=_CANCEL_POLL_SECONDS) + break + except subprocess.TimeoutExpired: + pass + if is_canceled is not None and is_canceled(): + raise CanceledException + if time.monotonic() > deadline: + raise SpliceError(f"ffmpeg timed out {description}") + finally: + if proc.returncode is None: + proc.kill() + proc.wait() + if proc.returncode != 0: + stderr.seek(0) + message = stderr.read().decode("utf-8", errors="replace").strip() + raise SpliceError(f"ffmpeg failed {description}: {message[-500:]}") + + +def copy_frames( + src: Path, + start: int, + end: int, + sample_count: int, + dst: Path, + is_canceled: Optional[Callable[[], bool]] = None, +) -> None: """Write ``dst`` = frames ``[start, end)`` of ``src``'s video, stream-copied. ``start`` and ``end`` are decode-order sample indices that must each be a keyframe (or @@ -373,6 +470,7 @@ def copy_frames(src: Path, start: int, end: int, sample_count: int, dst: Path) - run_ffmpeg( ["-i", str(src), "-map", "0:v:0", "-c", "copy", "-frames:v", str(end), "-f", "mp4", str(dst)], description, + is_canceled, ) return # The segment muxer splits at exact packet counts; the wanted range is the second segment. @@ -398,6 +496,7 @@ def copy_frames(src: Path, start: int, end: int, sample_count: int, dst: Path) - str(pattern), ], description, + is_canceled, ) segment = Path(work) / "segment1.mp4" if not segment.exists(): @@ -405,13 +504,28 @@ def copy_frames(src: Path, start: int, end: int, sample_count: int, dst: Path) - shutil.move(str(segment), str(dst)) +def probes_as(path: Path, fps: float, frames: int) -> bool: + """Whether a player's probe of ``path`` agrees with the timeline it should hold: ``frames`` + at ``fps``. A last check on a joined file, independent of the sample tables it was built + from — a squeezed or stretched timestamp shows up as a wrong rate or duration.""" + from invokeai.app.util.video_thumbnails import probe_video + + _, _, duration, probed_fps = probe_video(path) + # The probe reports duration to the hundredth of a second. + return ( + probed_fps is not None + and math.isclose(probed_fps, fps, rel_tol=1e-3) + and abs(duration - frames / fps) <= 0.5 / fps + 0.01 + ) + + def _concat_list_entry(path: Path) -> str: # The concat demuxer's quoting: close the quote, emit an escaped quote, reopen. escaped = str(path.resolve()).replace("'", "'\\''") return f"file '{escaped}'\n" -def concat_stream_copy(parts: list[Path], dst: Path) -> None: +def concat_stream_copy(parts: list[Path], dst: Path, is_canceled: Optional[Callable[[], bool]] = None) -> None: """Write ``dst`` = ``parts``' video streams joined end to end, stream-copied. Every part must share one decoder configuration (see :func:`read_h264_layout`) and @@ -439,6 +553,7 @@ def concat_stream_copy(parts: list[Path], dst: Path) -> None: str(dst), ], f"joining {len(parts)} stream-copied parts", + is_canceled, ) finally: list_path.unlink(missing_ok=True) diff --git a/invokeai/app/util/video_encoding.py b/invokeai/app/util/video_encoding.py index 0ae247fe746..91fcb1e80ee 100644 --- a/invokeai/app/util/video_encoding.py +++ b/invokeai/app/util/video_encoding.py @@ -71,11 +71,12 @@ def make_mp4_writer(path: Path | str, fps: float, audio_path: Path | str | None ) -_writer_decoder_configs: dict[tuple[int, int, float], bytes] = {} +_writer_sample_entrys: dict[tuple[int, int, float], bytes] = {} -def writer_decoder_config(width: int, height: int, fps: float) -> bytes | None: - """The ``avcC`` bytes ``make_mp4_writer`` writes for this canvas and rate. +def writer_sample_entry(width: int, height: int, fps: float) -> bytes | None: + """The sample entry ``make_mp4_writer`` writes for this canvas and rate (see + ``H264Layout.sample_entry``). They depend only on the encoder settings, canvas and rate, never on the content, so one encoded frame tells a splice up front whether a file can share a track with new output @@ -83,15 +84,15 @@ def writer_decoder_config(width: int, height: int, fps: float) -> bytes | None: successful probe is remembered. """ key = (width, height, fps) - if key not in _writer_decoder_configs: - config = _probe_decoder_config(width, height, fps) + if key not in _writer_sample_entrys: + config = _probe_sample_entry(width, height, fps) if config is None: return None - _writer_decoder_configs[key] = config - return _writer_decoder_configs[key] + _writer_sample_entrys[key] = config + return _writer_sample_entrys[key] -def _probe_decoder_config(width: int, height: int, fps: float) -> bytes | None: +def _probe_sample_entry(width: int, height: int, fps: float) -> bytes | None: from invokeai.app.util.h264_splice import read_h264_layout with tempfile.TemporaryDirectory(prefix="invokeai_encoder_probe_") as work: @@ -102,7 +103,7 @@ def _probe_decoder_config(width: int, height: int, fps: float) -> bytes | None: finally: writer.close() layout = read_h264_layout(path) - return layout.decoder_config if layout is not None else None + return layout.sample_entry if layout is not None else None def write_stereo_wav(path: Path | str, samples: np.ndarray, sample_rate: int) -> None: diff --git a/tests/app/invocations/test_video_concat_splice.py b/tests/app/invocations/test_video_concat_splice.py index 0f83d656c70..6336aa08263 100644 --- a/tests/app/invocations/test_video_concat_splice.py +++ b/tests/app/invocations/test_video_concat_splice.py @@ -22,8 +22,10 @@ import numpy as np import pytest +from invokeai.app.invocations import video_concat from invokeai.app.invocations.fields import VideoField from invokeai.app.invocations.video_concat import VideoConcatInvocation +from invokeai.app.services.session_processor.session_processor_common import CanceledException from invokeai.app.util.aac_splice import PREROLL_FRAMES from invokeai.app.util.h264_splice import read_h264_layout from invokeai.app.util.mp4_metadata import find_box, write_mp4_tags @@ -358,3 +360,74 @@ def test_spliced_soundtrack_matches_a_full_re_encode(tmp_path, monkeypatch, tran assert got is not None and want is not None n = min(got[0].shape[1], want[0].shape[1]) assert float(np.abs(got[0][:, :n] - want[0][:, :n]).max()) < 0.1 + + +@pytest.mark.parametrize( + "tag", + [("-color_primaries", "bt2020", "-color_trc", "smpte2084", "-colorspace", "bt2020nc"), ("-aspect", "16:9")], + ids=["colr", "pasp"], +) +def test_first_clip_with_presentation_tags_is_re_encoded(tmp_path, tag): + """The joined track would present the re-encoded continuation with the copied head's colour + or aspect tags, which were never applied to it.""" + _encode(tmp_path / "plain.mp4", _frames(100, seed=1)) + a = tmp_path / "a.mp4" + subprocess.run( + [imageio_ffmpeg.get_ffmpeg_exe(), "-y", "-loglevel", "error", "-i", str(tmp_path / "plain.mp4")] + + ["-c", "copy", *tag, str(a)], + check=True, + ) + b = _encode(tmp_path / "b.mp4", _frames(40, seed=2)) + out, context = _concat(tmp_path, [a, b], "out", transition="crossfade", transition_frames=2) + assert "Stream-copied" not in _logged(context) + layout = read_h264_layout(out) + assert layout is not None and b"colr" not in layout.sample_entry and b"pasp" not in layout.sample_entry + assert len(_decode(out)) == 138 + + +@pytest.mark.parametrize("stage", ["copy_frames", "mux_spliced_aac"]) +def test_cancellation_during_a_splice_saves_nothing(tmp_path, stage): + """Cancelling while a splice stage runs stops it — the stage itself raises, not just the + pre-save check — and nothing is saved.""" + a = _encode(tmp_path / "a.mp4", _frames(100, seed=1), tone=440.0) + b = _encode(tmp_path / "b.mp4", _frames(40, seed=2), tone=660.0) + canceled = [False] + raised: list[BaseException] = [] + real = getattr(video_concat, stage) + + def cancel_then_run(*args, **kwargs): + canceled[0] = True + try: + return real(*args, **kwargs) + except BaseException as e: + raised.append(e) + raise + + context = MagicMock() + context.videos.get_path.side_effect = lambda name: {"a": a, "b": b}[name] + context.util.is_canceled.side_effect = lambda: canceled[0] + node = VideoConcatInvocation(videos=[VideoField(video_name="a"), VideoField(video_name="b")], transition="cut") + with patch(f"invokeai.app.invocations.video_concat.{stage}", side_effect=cancel_then_run): + with pytest.raises(CanceledException): + node.invoke(context) + assert len(raised) == 1 and isinstance(raised[0], CanceledException) + context.videos.save.assert_not_called() + + +def test_tkhd_display_size_does_not_leak_onto_the_join(tmp_path): + """ffmpeg turns a track header display size that disagrees with the coded size into a + pasp when it copies the track: the copied head would then stretch the new clip too.""" + source = _encode(tmp_path / "plain.mp4", _frames(100, seed=1)) + data = bytearray(source.read_bytes()) + tkhd = data.find(b"tkhd") + assert data[tkhd + 4] == 0 # version 0: width at payload offset 76, as 16.16 fixed point + (width,) = struct.unpack_from(">I", data, tkhd + 4 + 76) + struct.pack_into(">I", data, tkhd + 4 + 76, width * 2) + a = tmp_path / "a.mp4" + a.write_bytes(bytes(data)) + b = _encode(tmp_path / "b.mp4", _frames(40, seed=2)) + + out, context = _concat(tmp_path, [a, b], "out", transition="cut", transition_frames=0) + layout = read_h264_layout(out) + assert layout is not None and b"pasp" not in layout.sample_entry + assert len(_decode(out)) == 140 diff --git a/tests/app/invocations/test_video_frame_extract_range.py b/tests/app/invocations/test_video_frame_extract_range.py index debbaa4f68f..1cd2dec8dfd 100644 --- a/tests/app/invocations/test_video_frame_extract_range.py +++ b/tests/app/invocations/test_video_frame_extract_range.py @@ -342,3 +342,31 @@ def test_squeezed_join_falls_back_to_re_encoding(tmp_path: Path, monkeypatch: py _, _, _, fps = probe_video(saved) assert fps == pytest.approx(16.0, rel=1e-3) assert len(_decode(saved)) == 65 + + +def test_cancellation_during_a_trim_splice_saves_nothing(tmp_path: Path) -> None: + source = tmp_path / "source.mp4" + _textured_source(source, 200) + canceled = [False] + raised: list[BaseException] = [] + from invokeai.app.invocations import video_frame_extract_range + + real = video_frame_extract_range.copy_frames + + def cancel_then_copy(*args, **kwargs): + canceled[0] = True + try: + return real(*args, **kwargs) + except BaseException as e: + raised.append(e) + raise + + context = MagicMock() + context.videos.get_path.return_value = source + context.util.is_canceled.side_effect = lambda: canceled[0] + invocation = ExtractVideoRangeInvocation(video=VideoField(video_name="source.mp4"), start_frame=31, end_frame=130) + with patch("invokeai.app.invocations.video_frame_extract_range.copy_frames", side_effect=cancel_then_copy): + with pytest.raises(CanceledException): + invocation.invoke(context) + assert len(raised) == 1 and isinstance(raised[0], CanceledException) + context.videos.save.assert_not_called() diff --git a/tests/app/util/test_h264_splice.py b/tests/app/util/test_h264_splice.py index 8612b155ade..4ec5191a796 100644 --- a/tests/app/util/test_h264_splice.py +++ b/tests/app/util/test_h264_splice.py @@ -6,16 +6,24 @@ by a stream copy; either would misalign or misorient a join. Real ffmpeg throughout. """ +import os +import struct import subprocess +import sys +import time +import tracemalloc from pathlib import Path +from unittest.mock import patch import imageio.v2 as iio2 import imageio_ffmpeg import numpy as np import pytest -from invokeai.app.util.h264_splice import read_h264_layout -from invokeai.app.util.video_encoding import make_mp4_writer, writer_decoder_config +from invokeai.app.services.session_processor.session_processor_common import CanceledException +from invokeai.app.util.aac_splice import read_aac_layout +from invokeai.app.util.h264_splice import read_h264_layout, run_ffmpeg +from invokeai.app.util.video_encoding import make_mp4_writer, write_stereo_wav, writer_sample_entry FPS = 16.0 @@ -52,7 +60,7 @@ def test_own_output_lists_exactly_the_decoded_frames(tmp_path: Path) -> None: assert layout is not None assert layout.sample_count == _decoded_count(path) == 70 assert layout.sync_samples[0] == 0 and len(layout.sync_samples) > 1 - assert layout.decoder_config == writer_decoder_config(96, 64, FPS) + assert layout.sample_entry == writer_sample_entry(96, 64, FPS) def test_edit_list_trim_is_refused(tmp_path: Path) -> None: @@ -101,3 +109,151 @@ def test_open_gop_i_frames_are_not_cut_points(tmp_path: Path) -> None: layout = read_h264_layout(path) assert layout is not None assert layout.sync_samples == (0,) + + +def _patch_box(data: bytearray, box: bytes, offset: int, value: bytes) -> None: + at = data.find(box) + assert at > 0, box + data[at + 4 + offset : at + 4 + offset + len(value)] = value + + +@pytest.mark.skipif(sys.platform == "win32", reason="bounds the child's address space with RLIMIT_AS") +@pytest.mark.parametrize("uniform_size", [0, 1], ids=["size-table", "uniform-size"]) +def test_oversized_sample_count_is_refused_without_allocating(tmp_path: Path, uniform_size: int) -> None: + """Uploads that are already H.264 MP4s are stored as-is, so the sample tables are the + uploader's. A count claiming ~4 billion samples, in a track with no sync-sample table + (so every sample would be a keyframe), must be refused before anything is sized by it.""" + data = bytearray(_encode(tmp_path / "a.mp4", 40).read_bytes()) + # stsz payload: version/flags, uniform sample size, sample count. + _patch_box(data, b"stsz", 4, struct.pack(">II", uniform_size, 0xFFFFFFF0)) + # Hide the sync-sample table: an all-sync track is the case that sized a tuple by the count. + at = data.find(b"stss") + assert at > 0 + data[at : at + 4] = b"free" + path = tmp_path / "hostile.mp4" + path.write_bytes(bytes(data)) + + script = ( + "import resource, sys\n" + # Far below the ~32 GB a count-sized tuple would take, with room for numpy's import. + "resource.setrlimit(resource.RLIMIT_AS, (4 << 30, 4 << 30))\n" + "from pathlib import Path\n" + "from invokeai.app.util.h264_splice import read_h264_layout\n" + "print(read_h264_layout(Path(sys.argv[1])))\n" + ) + result = subprocess.run( + [sys.executable, "-c", script, str(path)], + capture_output=True, + text=True, + timeout=120, + # One BLAS thread: each one reserves address space the limit would otherwise spend. + env={**os.environ, "OPENBLAS_NUM_THREADS": "1", "OMP_NUM_THREADS": "1"}, + ) + assert result.returncode == 0, result.stderr[-2000:] + assert result.stdout.strip() == "None" + + +@pytest.mark.parametrize( + "tag", + [ + ("-color_primaries", "bt2020", "-color_trc", "smpte2084", "-colorspace", "bt2020nc"), + ("-aspect", "16:9"), + ], + ids=["colr", "pasp"], +) +def test_presentation_tags_set_the_track_apart(tmp_path: Path, tag: tuple[str, ...]) -> None: + """A remux can add colour or pixel-aspect boxes while ``avcC`` stays byte-identical; a + joined track would present every frame with the copied part's tags.""" + source = _encode(tmp_path / "a.mp4", 40) + subprocess.run( + [imageio_ffmpeg.get_ffmpeg_exe(), "-y", "-loglevel", "error", "-i", str(source), "-c", "copy", *tag] + + [str(tmp_path / "tagged.mp4")], + check=True, + ) + plain, tagged = read_h264_layout(source), read_h264_layout(tmp_path / "tagged.mp4") + assert plain is not None and tagged is not None + assert tagged.decoder_config == plain.decoder_config + assert tagged.sample_entry != writer_sample_entry(96, 64, FPS) == plain.sample_entry + + +def test_cancellation_kills_a_running_ffmpeg(tmp_path: Path) -> None: + """A canceled node must not wait out a long copy: the child is killed and reaped.""" + started = time.monotonic() + calls = [0] + + def is_canceled() -> bool: + calls[0] += 1 + return time.monotonic() - started > 0.5 + + children: list[subprocess.Popen] = [] + real_popen = subprocess.Popen + + def spawn(*args, **kwargs) -> subprocess.Popen: + children.append(real_popen(*args, **kwargs)) + return children[-1] + + # An endless synthetic source that only stops when killed. + with patch("invokeai.app.util.h264_splice.subprocess.Popen", side_effect=spawn): + with pytest.raises(CanceledException): + run_ffmpeg(["-f", "lavfi", "-i", "anullsrc", "-f", "null", "-"], "running forever", is_canceled) + assert time.monotonic() - started < 10 + # The last child is the copy (locating the binary may spawn a version probe first). + assert children and children[-1].returncode is not None, "the child was left running" + + +def _grow_sample_entry(data: bytes, handler: bytes, blob: bytes) -> bytes: + """Append ``blob`` inside the ``handler`` track's sample entry, growing every ancestor box. + The file must keep ``moov`` after ``mdat`` so no chunk offset moves.""" + + def boxes(start: int, end: int): + while start + 8 <= end: + size, kind = struct.unpack(">I4s", data[start : start + 8]) + yield kind, start, size + start += size + + _, moov, moov_size = next(b for b in boxes(0, len(data)) if b[0] == b"moov") + for kind, trak, trak_size in boxes(moov + 8, moov + moov_size): + if kind != b"trak": + continue + chain = [(moov, moov_size), (trak, trak_size)] + start, end = trak + 8, trak + trak_size + for name in (b"mdia", b"minf", b"stbl", b"stsd"): + _, at, size = next(b for b in boxes(start, end) if b[0] == name) + chain.append((at, size)) + start, end = at + 8, at + size + hdlr = data.find(b"hdlr", trak, trak + trak_size) + if data[hdlr + 12 : hdlr + 16] != handler: + continue + _, entry, entry_size = next(boxes(start + 8, end)) + chain.append((entry, entry_size)) + grown = bytearray(data) + grown[entry + entry_size : entry + entry_size] = blob + for at, size in chain: + struct.pack_into(">I", grown, at, size + len(blob)) + return bytes(grown) + raise AssertionError(f"no {handler!r} track") + + +@pytest.mark.parametrize("handler", [b"vide", b"soun"]) +def test_sample_entry_stuffed_with_boxes_is_refused_cheaply(tmp_path: Path, handler: bytes) -> None: + """The sample entry is the uploader's: a flood of child boxes must be refused after a + bounded look, not collected (memory) or concatenated (quadratic time).""" + wav = tmp_path / "a.wav" + write_stereo_wav(wav, np.zeros((2, 48000), dtype=np.float32), 48000) + source = tmp_path / "a.mp4" + writer = make_mp4_writer(source, FPS, audio_path=wav) + for _ in range(16): + writer.append_data(np.zeros((64, 96, 3), dtype=np.uint8)) + writer.close() + hostile = tmp_path / "hostile.mp4" + hostile.write_bytes(_grow_sample_entry(source.read_bytes(), handler, struct.pack(">I4s", 8, b"free") * 200_000)) + + read = read_h264_layout if handler == b"vide" else read_aac_layout + tracemalloc.start() + try: + assert read(hostile) is None + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + # 200k boxes are 1.6 MB of file; reading a bounded handful of them allocates almost nothing. + assert peak < 1 << 20 From de28b6aedd508a8df6d71bf9aea131767a94c526 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Mon, 12 Oct 2026 22:17:05 -0400 Subject: [PATCH 3/7] fix(video): stop re-encoding the soundtrack kept by a trim extract_video_range stream-copies the source's AAC packets inside the range, encoding only a lead-in and tail; video_concat shares the same verified splice. Soundtracks are encoded without PNS so copied packets decode exactly after a lead-in. --- .../05.Video Generation/introduction.mdx | 2 +- invokeai/app/invocations/video_concat.py | 89 ++---- .../invocations/video_frame_extract_range.py | 89 +++++- invokeai/app/util/aac_splice.py | 268 ++++++++++++------ invokeai/app/util/video_audio.py | 6 +- invokeai/app/util/video_encoding.py | 6 +- invokeai/app/util/video_ingest.py | 4 +- .../invocations/test_video_concat_splice.py | 2 +- ..._video_frame_extract_range_audio_splice.py | 222 +++++++++++++++ 9 files changed, 531 insertions(+), 157 deletions(-) create mode 100644 tests/app/invocations/test_video_frame_extract_range_audio_splice.py diff --git a/docs/src/content/docs/Users Guide/05.Video Generation/introduction.mdx b/docs/src/content/docs/Users Guide/05.Video Generation/introduction.mdx index 3cf8efa3d2b..5b870c8039a 100644 --- a/docs/src/content/docs/Users Guide/05.Video Generation/introduction.mdx +++ b/docs/src/content/docs/Users Guide/05.Video Generation/introduction.mdx @@ -203,7 +203,7 @@ Extension is available on all three families. LTX-2's continues the soundtrack a ### Quality drift across iterations -Each iteration's source is itself a generation output, so artifacts compound — by the 4th or 5th extension you may see softening or color drift. That drift is in the *new* footage: the segments already in the initial video are carried into the result without being re-encoded, picture and sound, apart from up to two seconds of picture at each cut — before the join, and at a trim's start. A trimmed initial video's soundtrack is re-encoded once. Extending a clip whose frame rate differs from the model's (MiniMax H3 always runs at 24 fps) re-encodes it in full, once, as does extending an uploaded video or one made by an earlier version of Invoke. Mitigations: +Each iteration's source is itself a generation output, so artifacts compound — by the 4th or 5th extension you may see softening or color drift. That drift is in the *new* footage: the segments already in the initial video are carried into the result without being re-encoded, picture and sound, apart from up to two seconds of picture and about a tenth of a second of sound at each cut — before the join, and at each end of a trim. Extending a clip whose frame rate differs from the model's (MiniMax H3 always runs at 24 fps) re-encodes it in full, once, as does extending an uploaded video or one made by an earlier version of Invoke. Mitigations: 1. **Trim the bridge point a few frames back from the end** — boundary frames are the most artifact-prone. 2. **Refresh a bridge frame with a low-strength img2img pass** (SDXL or FLUX at ~0.2 strength) before using it as a first frame for the next segment. diff --git a/invokeai/app/invocations/video_concat.py b/invokeai/app/invocations/video_concat.py index 0701a488b09..cb4dfe5126b 100644 --- a/invokeai/app/invocations/video_concat.py +++ b/invokeai/app/invocations/video_concat.py @@ -52,10 +52,10 @@ from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.app.util.aac_splice import ( AAC_FRAME_SAMPLES, + EDGE_FRAMES, PREROLL_FRAMES, - encoder_aac_layout, - mux_spliced_aac, - read_aac_layout, + splice_soundtrack, + spliceable_aac_layout, ) from invokeai.app.util.h264_splice import ( MIN_ENCODED_FRAMES, @@ -67,7 +67,6 @@ read_h264_layout, ) from invokeai.app.util.video_audio import ( - AudioExtractionError, extract_audio_pcm, mux_audio_into_video, resample_linear, @@ -455,15 +454,12 @@ def _mux_spliced_audio( transition — provided that clip plays at the output rate, so its audio is not retimed, and sets the soundtrack's sample rate. Returns False, leaving the caller to encode the whole soundtrack, when the clip's track is not one this app's encoder - writes, nothing whole-frame precedes the transition, or the result's head does not - decode back to exactly the samples it should hold. + writes, too little precedes the transition, or the splice does not verify. """ if first_fps is None or first_fps != output_fps: return False - layout, encoder = read_aac_layout(first_path), encoder_aac_layout(rate) - if layout is None or encoder is None or layout.sample_rate != rate: - return False - if (layout.decoder_config, layout.priming) != (encoder.decoder_config, encoder.priming): + layout = spliceable_aac_layout(first_path, rate) + if layout is None: return False # The soundtrack holds the clip's own decoded samples only where its audio was mapped # one-to-one: the clip's span at its native rate must be exactly its span on the timeline. @@ -475,66 +471,31 @@ def _mux_spliced_audio( layout.packet_count * AAC_FRAME_SAMPLES - layout.priming, pcm.shape[1], ) - preroll = PREROLL_FRAMES * AAC_FRAME_SAMPLES - # Keep the cut PREROLL_FRAMES clear of the first changed sample (see aac_splice). - head_packets = (unchanged + layout.priming) // AAC_FRAME_SAMPLES - PREROLL_FRAMES - cut = head_packets * AAC_FRAME_SAMPLES - layout.priming - if cut < preroll: - return False - try: - seam_reference = mux_spliced_aac( - video_path, - first_path, - head_packets, - layout.priming, - pcm[:, cut - preroll :], - rate, - pcm.shape[1], - out_path, - context.util.is_canceled, - ) - muxed = extract_audio_pcm(out_path) - except (SpliceError, AudioExtractionError) as e: - context.logger.info(f"Re-encoding the whole soundtrack: it cannot be stream-copied ({e})") - return False - reason = self._spliced_audio_defect(muxed, pcm, rate, cut, seam_reference) + # Keep the cut PREROLL_FRAMES clear of the first changed sample, and a full encoded tail + # after it (see aac_splice). + head_packets = min( + (unchanged + layout.priming) // AAC_FRAME_SAMPLES - PREROLL_FRAMES, + (pcm.shape[1] - EDGE_FRAMES * AAC_FRAME_SAMPLES + layout.priming) // AAC_FRAME_SAMPLES, + ) + reason = splice_soundtrack( + video_path, + first_path, + layout.priming, + 0, + head_packets, + -layout.priming, + pcm, + rate, + out_path, + context.util.is_canceled, + ) if reason is not None: context.logger.info(f"Re-encoding the whole soundtrack: {reason}") return False + cut = head_packets * AAC_FRAME_SAMPLES - layout.priming context.logger.info(f"Stream-copied the first {cut} audio samples of the first clip without re-encoding") return True - @staticmethod - def _spliced_audio_defect( - muxed: Optional[tuple[np.ndarray, int]], - pcm: np.ndarray, - rate: int, - cut: int, - seam_reference: np.ndarray, - ) -> Optional[str]: - """Why the spliced soundtrack cannot be used, or None. - - The copied head must decode back to exactly the samples it stands for, and the - seam — the frames where copied and new packets overlap — must be no noisier than an - ordinary encode of the same samples (``seam_reference``). Whether the overlap - cancels depends on the encoder's window decisions around the content there, so it is - measured, not assumed; and it is measured against the seam's own content, since the - audio after it may be silence that codes with no error at all. - """ - if muxed is None or muxed[1] != rate or muxed[0].shape[1] < pcm.shape[1]: - return "the spliced soundtrack did not decode to its full length" - decoded = muxed[0] - if not np.array_equal(decoded[:, :cut], pcm[:, :cut]): - return "the copied audio did not decode back exactly" - seam = slice(cut, cut + seam_reference.shape[1]) - if seam_reference.shape[1] != PREROLL_FRAMES * AAC_FRAME_SAMPLES or seam.stop > pcm.shape[1]: - return "the new audio does not cover the seam" - seam_error = float(np.sqrt(np.mean((decoded[:, seam] - pcm[:, seam]) ** 2))) - coding_error = float(np.sqrt(np.mean((seam_reference - pcm[:, seam]) ** 2))) - if seam_error > 2.0 * coding_error + 1e-3: - return "the seam between copied and new audio is not clean" - return None - def _estimate_transition_memory(self, width: int, height: int) -> int: if self.transition == "cut" or self.transition_frames == 0: return 0 diff --git a/invokeai/app/invocations/video_frame_extract_range.py b/invokeai/app/invocations/video_frame_extract_range.py index fd798478c8b..b71a9e919b8 100644 --- a/invokeai/app/invocations/video_frame_extract_range.py +++ b/invokeai/app/invocations/video_frame_extract_range.py @@ -11,7 +11,9 @@ through this node (as video extension does with an untrimmed source) costs no quality. A trim of a source this app encoded stream-copies the keyframe-aligned run of frames inside the range and encodes only the partial keyframe intervals -at either end (see ``h264_splice``). Anything else is re-encoded in full. +at either end (see ``h264_splice``), and likewise copies the source's AAC +packets for the audio inside the range, encoding only a short lead-in and tail +(see ``aac_splice``). Anything else is re-encoded in full. If the source carries an audio track, the matching slice of it is carried through (and retimed alongside the video when the output fps differs from @@ -48,6 +50,12 @@ from invokeai.app.invocations.primitives import VideoOutput from invokeai.app.services.session_processor.session_processor_common import CanceledException from invokeai.app.services.shared.invocation_context import InvocationContext +from invokeai.app.util.aac_splice import ( + AAC_FRAME_SAMPLES, + EDGE_FRAMES, + splice_soundtrack, + spliceable_aac_layout, +) from invokeai.app.util.h264_splice import ( MIN_ENCODED_FRAMES, H264Layout, @@ -281,14 +289,17 @@ def invoke(self, context: InvocationContext) -> ExtractVideoRangeOutput: if audio is not None: pcm, rate = audio context.util.signal_progress("Muxing audio") - wav_tmp = tempfile.NamedTemporaryFile(prefix="invokeai_video_range_", suffix=".wav", delete=False) - wav_tmp.close() - wav_path = Path(wav_tmp.name) - write_stereo_wav(wav_path, pcm, rate) mux_tmp = tempfile.NamedTemporaryFile(prefix="invokeai_video_range_", suffix=".mp4", delete=False) mux_tmp.close() muxed_path = Path(mux_tmp.name) - mux_audio_into_video(tmp_path, wav_path, muxed_path) + if not self._mux_spliced_audio( + context, video_path, start, end, source_fps, output_fps, pcm, rate, tmp_path, muxed_path + ): + wav_tmp = tempfile.NamedTemporaryFile(prefix="invokeai_video_range_", suffix=".wav", delete=False) + wav_tmp.close() + wav_path = Path(wav_tmp.name) + write_stereo_wav(wav_path, pcm, rate) + mux_audio_into_video(tmp_path, wav_path, muxed_path) source_path = muxed_path out_duration = num_frames / output_fps @@ -466,6 +477,68 @@ def _save_unchanged( end_frame=n_frames - 1, ) + def _mux_spliced_audio( + self, + context: InvocationContext, + video_path: Path, + start: int, + end: int, + source_fps: Optional[float], + output_fps: float, + pcm: np.ndarray, + rate: int, + trimmed_video: Path, + out_path: Path, + ) -> bool: + """Mux ``pcm`` onto ``trimmed_video``, stream-copying the source's own AAC packets + for the part of the range they hold whole (see ``aac_splice``). + + ``pcm`` is the source's decoded audio over the range, sample for sample, when the + range plays at the source's own rate and its span maps one-to-one onto the output + timeline; the source's packets are then copied for every output frame that lies + inside both the range and the source's real audio, and only the partial frames at + each end are encoded. Returns False, leaving the caller to encode the whole + soundtrack, when the source's track is not one this app's encoder writes, the range + is too short, or the splice does not verify. + """ + if source_fps is None or source_fps != output_fps: + return False + window_start = round(start * rate / source_fps) + window_end = round((end + 1) * rate / source_fps) + if window_end - window_start != pcm.shape[1]: + return False + layout = spliceable_aac_layout(video_path, rate) + if layout is None: + return False + frame, priming = AAC_FRAME_SAMPLES, layout.priming + # Packet k's output frame is decoded samples [k * frame - priming, (k + 1) * frame - priming). + if window_start == 0: + first, copied_at = 0, -priming + else: + # The first frame wholly inside the range that leaves a full lead-in before it. + first = -(-(window_start + priming + EDGE_FRAMES * frame) // frame) + copied_at = first * frame - priming - window_start + # The last packet that leaves a full encoded tail after the run (see aac_splice); past the + # source's own audio the tail is the silence the timeline was padded with. + end_packet = min(layout.packet_count, (window_end - EDGE_FRAMES * frame + priming) // frame) + reason = splice_soundtrack( + trimmed_video, + video_path, + priming, + first, + end_packet, + copied_at, + pcm, + rate, + out_path, + context.util.is_canceled, + ) + if reason is not None: + context.logger.info(f"Re-encoding the whole soundtrack: {reason}") + return False + context.logger.info(f"Stream-copied audio packets {first}-{end_packet - 1} of the source without re-encoding") + return True + def _build_audio_track( self, context: InvocationContext, @@ -511,7 +584,9 @@ def _build_audio_track( # Audio track shorter than the video (or the range lies past its end): # keep temporal alignment by padding the shortfall with silence. window = np.pad(window, ((0, 0), (0, (window_end - window_start) - window.shape[1]))) - target = round((end - start + 1) / output_fps * rate) + # At the source's own rate the window maps sample for sample; rounding the span on its + # own could come out one off the window and interpolate the whole of it. + target = window.shape[1] if output_fps == eff_fps else round((end - start + 1) / output_fps * rate) return resample_linear(window, target), rate @staticmethod diff --git a/invokeai/app/util/aac_splice.py b/invokeai/app/util/aac_splice.py index ddb8b412219..e31d140a029 100644 --- a/invokeai/app/util/aac_splice.py +++ b/invokeai/app/util/aac_splice.py @@ -10,9 +10,11 @@ packet after them blends with the last copied one. That blend cancels cleanly only when both packets were coded from the same signal across their shared window, so two things keep the seam at coding-noise level rather than a glitch: the new encode starts -``PREROLL_FRAMES`` frames before the cut (one is not enough — the encoder's first frames -after a cold start make different window decisions), and callers keep the cut that many -frames clear of anything the joined soundtrack changes. Callers also verify the seam. +``PREROLL_FRAMES`` frames before the cut (the encoder takes several frames after a cold +start to settle its window and bit-allocation decisions: one frame of pre-roll leaves +glitches, two still leaves an occasional 40 ms burst on music), and callers keep the cut +that many frames clear of anything the joined soundtrack changes. Callers also verify the +seam. Packets are joined as ADTS, which carries no timing, so two things the MP4 encodes in its edit list are restored afterwards: the encoder priming the source's first packet @@ -34,12 +36,16 @@ from invokeai.app.util.h264_splice import SpliceError, run_ffmpeg from invokeai.app.util.mp4_metadata import find_box, iter_boxes -from invokeai.app.util.video_audio import AAC_ENCODE_ARGS, extract_audio_pcm +from invokeai.app.util.video_audio import AAC_ENCODE_ARGS, AudioExtractionError, extract_audio_pcm AAC_FRAME_SAMPLES = 1024 # Frames of already-known audio the new encode starts with before the cut, and the clearance # the cut keeps from the first changed sample. -PREROLL_FRAMES = 2 +PREROLL_FRAMES = 4 +# The fewest whole frames an encoded lead-in or tail holds beside the copied run: the frame +# it shares with the run plus PREROLL_FRAMES more, so the seam is far from the encoder's cold +# start (a lead-in) or its end-of-stream flush (a tail). +EDGE_FRAMES = PREROLL_FRAMES + 1 _MAX_TOP_LEVEL_BOXES = 1024 _MAX_PACKETS = 1 << 24 # AudioSampleEntry fields before its child boxes (SampleEntry 8 + version-0 audio header 20). @@ -57,7 +63,9 @@ class AacLayout(NamedTuple): sample_rate: int """The media timescale, which for AAC is the sample rate.""" priming: int - """Leading samples the edit list skips: the encoder delay held by the first packet.""" + """Leading samples the edit list skips. For a file the encoder wrote directly this is its + priming; a splice that opens on a lead-in skips the lead-in's padding too. Either way + packet ``k`` decodes to presented samples ``[k * 1024 - priming, (k + 1) * 1024 - priming)``.""" decoder_config: bytes """The ``AudioSpecificConfig`` (object type, sample rate, channels), without the backward-compatible "no SBR" extension ffmpeg appends to its own encodes but cannot @@ -232,6 +240,18 @@ def _probe_encoder_layout(sample_rate: int) -> Optional[AacLayout]: return read_aac_layout(probe) +def spliceable_aac_layout(path: Path, sample_rate: int) -> Optional[AacLayout]: + """``path``'s audio layout if its packets can be spliced into a soundtrack this app + encodes at ``sample_rate``: the same AudioSpecificConfig as its encoder, and a leading edit + that skips at least that encoder's priming (anything less would be another encoder's).""" + layout, encoder = read_aac_layout(path), encoder_aac_layout(sample_rate) + if layout is None or encoder is None or layout.sample_rate != sample_rate: + return None + if layout.decoder_config != encoder.decoder_config or layout.priming < encoder.priming: + return None + return layout + + def _adts_frames(data: bytes) -> list[bytes]: frames: list[bytes] = [] position = 0 @@ -247,9 +267,11 @@ def _adts_frames(data: bytes) -> list[bytes]: return frames -def _set_track_duration(path: Path, handler: bytes, samples: int, sample_rate: int) -> None: - """Trim the ``handler`` track's presentation to ``samples`` via its edit list, and keep - ``tkhd`` and ``mvhd`` durations consistent with it. All fields are rewritten in place.""" +def _set_audio_edit(path: Path, media_time: int, samples: int, sample_rate: int) -> None: + """Make the audio track present ``samples`` samples starting ``media_time`` samples into + its media, via its edit list, and keep ``tkhd`` and ``mvhd`` durations consistent with it. + All fields are rewritten in place.""" + handler = b"soun" with open(path, "r+b") as fh: fh.seek(0, 2) moov = find_box(fh, 0, fh.tell(), b"moov", max_boxes=_MAX_TOP_LEVEL_BOXES) @@ -282,9 +304,9 @@ def _set_track_duration(path: Path, handler: bytes, samples: int, sample_rate: i if elst is None: raise SpliceError(f"{path.name}'s {handler.decode()} track has no edit list") fh.seek(elst[0]) - elst_field = ">Q" if fh.read(1)[0] == 1 else ">I" + wide_edit = fh.read(1)[0] == 1 fh.seek(elst[0] + 8) - fh.write(struct.pack(elst_field, duration)) + fh.write(struct.pack(">Qq" if wide_edit else ">Ii", duration, media_time)) fh.seek(tkhd_duration_at) fh.write(struct.pack(field, duration)) patched = True @@ -296,79 +318,163 @@ def _set_track_duration(path: Path, handler: bytes, samples: int, sample_rate: i fh.write(struct.pack(">Q" if mvhd_version == 1 else ">I", max(track_durations))) -def mux_spliced_aac( +def _encode( + pcm: np.ndarray, sample_rate: int, work: Path, name: str, is_canceled: Optional[Callable[[], bool]] +) -> tuple[list[bytes], np.ndarray]: + """Encode ``pcm`` as the app does; returns its ADTS frames, priming frames included, and + their standalone decode (which, knowing nothing of priming, starts with the priming).""" + from invokeai.app.util.video_encoding import write_stereo_wav + + wav, adts = work / f"{name}.wav", work / f"{name}.aac" + write_stereo_wav(wav, pcm, sample_rate) + run_ffmpeg(["-i", str(wav), *AAC_ENCODE_ARGS, "-f", "adts", str(adts)], "encoding the new audio", is_canceled) + decoded = extract_audio_pcm(adts) + if decoded is None: + raise SpliceError("the new audio decoded to nothing") + return _adts_frames(adts.read_bytes()), decoded[0] + + +def _seam_defect(decoded: np.ndarray, timeline: np.ndarray, start: int, reference: np.ndarray) -> Optional[str]: + """Whether the seam ``decoded[:, start:]`` is noisier than an ordinary encode of the same + samples (``reference``). Whether copied and new packets overlap cleanly depends on the + encoder's window decisions around the content there, so it is measured, not assumed — and + against the seam's own content, since the audio either side may be silence that codes with + no error at all.""" + length = min(reference.shape[1], timeline.shape[1] - start) + if start < 0 or length <= 0: + return "the new audio does not cover a seam" + seam, reference = slice(start, start + length), reference[:, :length] + seam_error = float(np.sqrt(np.mean((decoded[:, seam] - timeline[:, seam]) ** 2))) + coding_error = float(np.sqrt(np.mean((reference - timeline[:, seam]) ** 2))) + if seam_error > 2.0 * coding_error + 1e-3: + return "a seam between copied and new audio is not clean" + return None + + +def splice_soundtrack( video_path: Path, source: Path, - head_packets: int, - priming: int, - tail_pcm: np.ndarray, + source_skip: int, + first: int, + end: int, + copied_at: int, + timeline: np.ndarray, sample_rate: int, - total_samples: int, out_path: Path, is_canceled: Optional[Callable[[], bool]] = None, -) -> np.ndarray: - """Write ``out_path`` = ``video_path``'s video plus a soundtrack spliced from ``source``. - - The soundtrack is ``source``'s first ``head_packets`` AAC packets, stream-copied — - ``head_packets * 1024 - priming`` samples after priming — followed by ``tail_pcm`` - encoded from ``PREROLL_FRAMES`` frames *before* that cut (``tail_pcm`` must start - there), presented for ``total_samples`` samples in all. - - Returns the new encode's own decode of the ``PREROLL_FRAMES`` frames after the cut: - the coding noise this content gets from an ordinary encode, which is what a clean - seam over the same samples should be measured against. +) -> Optional[str]: + """Write ``out_path`` = ``video_path``'s video with ``timeline`` as its soundtrack, + stream-copying ``source``'s AAC packets ``[first, end)`` for the part of it they hold. + + ``timeline`` (``(2, n)`` float PCM) is the soundtrack to present; packet ``first``'s + output frame begins at ``timeline[:, copied_at]``, and the copied packets must decode to + the timeline's samples there. ``source_skip`` is the source's leading edit (its layout's + ``priming``). ``first == 0`` copies the source's own start, leading edit included + (``copied_at == -source_skip``); otherwise the frames before the copied run are + encoded as a lead-in. The frames after it are encoded from ``PREROLL_FRAMES`` frames + before its end. Both seams are verified, as is every copied frame not on a seam, which + must decode bit-exact. + + Returns None on success, or why the soundtrack could not be spliced — ``out_path`` is + then unusable and the caller encodes the whole soundtrack instead. """ - from invokeai.app.util.video_encoding import write_stereo_wav - - if priming % AAC_FRAME_SAMPLES: - raise SpliceError(f"priming of {priming} samples is not a whole AAC frame") - with tempfile.TemporaryDirectory(prefix="invokeai_aac_splice_") as work: - head, tail_wav, tail = Path(work) / "head.aac", Path(work) / "tail.wav", Path(work) / "tail.aac" - run_ffmpeg( - ["-i", str(source), "-map", "0:a:0", "-c", "copy", "-frames:a", str(head_packets), "-f", "adts", str(head)], - f"copying the first {head_packets} audio packets of {source.name}", - is_canceled, - ) - write_stereo_wav(tail_wav, tail_pcm, sample_rate) - run_ffmpeg( - ["-i", str(tail_wav), *AAC_ENCODE_ARGS, "-f", "adts", str(tail)], "encoding the new audio", is_canceled - ) - head_frames = _adts_frames(head.read_bytes()) - if len(head_frames) != head_packets: - raise SpliceError(f"copied {len(head_frames)} audio packets, expected {head_packets}") - # Drop the new encode's priming packets and its pre-roll. - dropped = priming // AAC_FRAME_SAMPLES + PREROLL_FRAMES - tail_frames = _adts_frames(tail.read_bytes())[dropped:] - decoded = extract_audio_pcm(tail) - if decoded is None: - raise SpliceError("the new audio decoded to nothing") - # An ADTS decode knows nothing of priming, so it starts with the priming frame. - seam_start = priming + PREROLL_FRAMES * AAC_FRAME_SAMPLES - seam_reference = decoded[0][:, seam_start : seam_start + PREROLL_FRAMES * AAC_FRAME_SAMPLES] - joined = Path(work) / "joined.aac" - joined.write_bytes(b"".join(head_frames + tail_frames)) - run_ffmpeg( - [ - "-i", - str(video_path), - # The first copied packet still holds the source's priming; presenting it from - # -priming makes ffmpeg write the leading edit that skips it. - "-itsoffset", - repr(-priming / sample_rate), - "-i", - str(joined), - "-map", - "0:v:0", - "-map", - "1:a:0", - "-c", - "copy", - "-f", - "mp4", - str(out_path), - ], - "muxing the spliced audio", - is_canceled, - ) - _set_track_duration(out_path, b"soun", total_samples, sample_rate) - return seam_reference + encoder = encoder_aac_layout(sample_rate) + if encoder is None or encoder.priming % AAC_FRAME_SAMPLES: + return "the AAC encoder's priming is not a whole number of frames" + frame, priming = AAC_FRAME_SAMPLES, encoder.priming + priming_frames = priming // frame + copied_end = copied_at + (end - first) * frame + if (first == 0) != (copied_at == -source_skip) or end - first < PREROLL_FRAMES: + return "the copied packets do not fit the soundtrack" + if first > 0 and copied_at < EDGE_FRAMES * frame: + return "the copied packets do not fit the soundtrack" + if copied_end < timeline.shape[1] and timeline.shape[1] - copied_end < EDGE_FRAMES * frame: + return "the copied packets do not fit the soundtrack" + if copied_end > timeline.shape[1] or copied_end - PREROLL_FRAMES * frame < 0: + return "the copied packets do not fit the soundtrack" + try: + with tempfile.TemporaryDirectory(prefix="invokeai_aac_splice_") as work: + work_dir = Path(work) + # The source's packets as ADTS, sliced by index: an ADTS frame per packet. + source_adts = work_dir / "source.aac" + run_ffmpeg( + ["-i", str(source), "-map", "0:a:0", "-c", "copy", "-f", "adts", str(source_adts)], + f"copying the audio packets of {source.name}", + is_canceled, + ) + source_frames = _adts_frames(source_adts.read_bytes()) + if len(source_frames) < end: + return f"{source.name} holds {len(source_frames)} audio packets, not {end}" + frames: list[bytes] = [] + seams: list[tuple[int, np.ndarray]] = [] + exact_from = 0 + if first == 0: + # The stream opens with the source's own first packet, and keeps its leading edit. + media_time = source_skip + else: + # A lead-in encode up to the copied run, padded at its start so its frames fall on + # the copied packets' grid, and run on past the seam so its last kept packet was + # coded with the real signal in the window it shares with the first copied one. + pad = (-copied_at) % frame + lead_frames = (pad + copied_at) // frame + lead_input = np.concatenate( + [ + np.zeros((2, pad), dtype=np.float32), + timeline[:, : min(timeline.shape[1], copied_at + PREROLL_FRAMES * frame)], + ], + axis=1, + ) + lead, lead_decoded = _encode(lead_input, sample_rate, work_dir, "lead", is_canceled) + frames += lead[: priming_frames + lead_frames] + media_time = priming + pad + seam_at = copied_at - frame + seams.append((seam_at, lead_decoded[:, priming + pad + seam_at :][:, : 2 * frame])) + exact_from = copied_at + frame + frames += source_frames[first:end] + if copied_end < timeline.shape[1]: + tail, tail_decoded = _encode( + timeline[:, copied_end - PREROLL_FRAMES * frame :], sample_rate, work_dir, "tail", is_canceled + ) + frames += tail[priming_frames + PREROLL_FRAMES :] + seams.append((copied_end, tail_decoded[:, priming + PREROLL_FRAMES * frame :][:, : 2 * frame])) + joined = work_dir / "joined.aac" + joined.write_bytes(b"".join(frames)) + run_ffmpeg( + [ + "-i", + str(video_path), + # Presenting the audio from -media_time makes ffmpeg write a leading edit, whose + # fields are then set exactly below. + "-itsoffset", + repr(-media_time / sample_rate), + "-i", + str(joined), + "-map", + "0:v:0", + "-map", + "1:a:0", + "-c", + "copy", + "-f", + "mp4", + str(out_path), + ], + "muxing the spliced audio", + is_canceled, + ) + _set_audio_edit(out_path, media_time, timeline.shape[1], sample_rate) + muxed = extract_audio_pcm(out_path) + except (SpliceError, AudioExtractionError, OSError) as e: + return f"it cannot be stream-copied ({e})" + + if muxed is None or muxed[1] != sample_rate or muxed[0].shape[1] < timeline.shape[1]: + return "the spliced soundtrack did not decode to its full length" + decoded = muxed[0] + interior = slice(max(exact_from, 0), copied_end) + if not np.array_equal(decoded[:, interior], timeline[:, interior]): + return "the copied audio did not decode back exactly" + for start, reference in seams: + defect = _seam_defect(decoded, timeline, start, reference) + if defect is not None: + return defect + return None diff --git a/invokeai/app/util/video_audio.py b/invokeai/app/util/video_audio.py index cd7ed1d3fbf..eca1db28281 100644 --- a/invokeai/app/util/video_audio.py +++ b/invokeai/app/util/video_audio.py @@ -24,7 +24,11 @@ import numpy as np # How every soundtrack the app writes is encoded; ``aac_splice`` matches a spliced tail to it. -AAC_ENCODE_ARGS = ["-c:a", "aac"] +# Perceptual noise substitution is off: the decoder fills those bands from a random generator +# whose state carries across packets, so a freshly encoded lead-in that used it would change +# how the stream-copied packets after it decode, and a splice could not verify them exact. +AAC_ENCODER_OPTIONS = ["-aac_pns", "0"] +AAC_ENCODE_ARGS = ["-c:a", "aac", *AAC_ENCODER_OPTIONS] def _ffmpeg_exe() -> str: diff --git a/invokeai/app/util/video_encoding.py b/invokeai/app/util/video_encoding.py index 91fcb1e80ee..906fcf590f1 100644 --- a/invokeai/app/util/video_encoding.py +++ b/invokeai/app/util/video_encoding.py @@ -36,6 +36,8 @@ import imageio.v2 as iio2 import numpy as np +from invokeai.app.util.video_audio import AAC_ENCODER_OPTIONS + CRF = 18 KEYFRAME_INTERVAL_SECONDS = 2.0 @@ -57,6 +59,8 @@ def make_mp4_writer(path: Path | str, fps: float, audio_path: Path | str | None raise FileNotFoundError(f"Audio file for muxing does not exist: {audio_path}") kwargs["audio_path"] = str(audio_path) kwargs["audio_codec"] = audio_codec + # Soundtracks share the WAV-mux path's encoder options so either can be spliced onto the other. + audio_options = AAC_ENCODER_OPTIONS if audio_path is not None and audio_codec == "aac" else [] keyframe_interval = max(1, round(fps * KEYFRAME_INTERVAL_SECONDS)) return iio2.get_writer( str(path), @@ -66,7 +70,7 @@ def make_mp4_writer(path: Path | str, fps: float, audio_path: Path | str | None codec="libx264", macro_block_size=1, quality=None, - output_params=["-crf", str(CRF), "-g", str(keyframe_interval)], + output_params=["-crf", str(CRF), "-g", str(keyframe_interval), *audio_options], **kwargs, ) diff --git a/invokeai/app/util/video_ingest.py b/invokeai/app/util/video_ingest.py index 29b51c4c35a..a37a8c8e571 100644 --- a/invokeai/app/util/video_ingest.py +++ b/invokeai/app/util/video_ingest.py @@ -35,6 +35,8 @@ from pathlib import Path from typing import Literal, Optional +from invokeai.app.util.video_audio import AAC_ENCODER_OPTIONS + # Hard wall-clock bound for a single ffmpeg child (probe or convert). Matches the bound # used for audio extraction: generous enough for a worst-case ~1 GB 4K transcode on a # modest CPU, small enough that a hostile file cannot camp an upload slot forever. @@ -123,7 +125,7 @@ def _audio_args(probe: MediaProbe) -> list[str]: return [] if probe.audio_codec == "aac": return ["-c:a", "copy"] - return ["-c:a", "aac", "-b:a", "192k"] + return ["-c:a", "aac", "-b:a", "192k", *AAC_ENCODER_OPTIONS] def plan_ingest(probe: MediaProbe) -> IngestAction: diff --git a/tests/app/invocations/test_video_concat_splice.py b/tests/app/invocations/test_video_concat_splice.py index 6336aa08263..b108ff82d6b 100644 --- a/tests/app/invocations/test_video_concat_splice.py +++ b/tests/app/invocations/test_video_concat_splice.py @@ -385,7 +385,7 @@ def test_first_clip_with_presentation_tags_is_re_encoded(tmp_path, tag): assert len(_decode(out)) == 138 -@pytest.mark.parametrize("stage", ["copy_frames", "mux_spliced_aac"]) +@pytest.mark.parametrize("stage", ["copy_frames", "splice_soundtrack"]) def test_cancellation_during_a_splice_saves_nothing(tmp_path, stage): """Cancelling while a splice stage runs stops it — the stage itself raises, not just the pre-save check — and nothing is saved.""" diff --git a/tests/app/invocations/test_video_frame_extract_range_audio_splice.py b/tests/app/invocations/test_video_frame_extract_range_audio_splice.py new file mode 100644 index 00000000000..d33d649ab8f --- /dev/null +++ b/tests/app/invocations/test_video_frame_extract_range_audio_splice.py @@ -0,0 +1,222 @@ +"""A trimmed source's soundtrack must not be re-encoded where the trim keeps it. + +Extend Video trims its source before joining a continuation onto it. The trim used to +decode and re-encode the whole kept soundtrack, so every extension of a trimmed source +cost the earlier audio one more AAC generation even though its video was stream-copied. +The node now copies the source's AAC packets for every output frame wholly inside the +range and encodes only a lead-in and a tail. These tests pin that the copied audio comes +through bit-exact, that the seams are no worse than a full re-encode, that the soundtrack +ends with the video, and that a source the splice cannot vouch for is re-encoded instead. + +Real ffmpeg (the imageio-ffmpeg bundled binary) encodes and decodes throughout. +""" + +import struct +from pathlib import Path +from unittest.mock import MagicMock, patch + +import imageio.v2 as iio2 +import numpy as np +import pytest + +from invokeai.app.invocations.fields import VideoField +from invokeai.app.invocations.video_concat import VideoConcatInvocation +from invokeai.app.invocations.video_frame_extract_range import ExtractVideoRangeInvocation +from invokeai.app.util.aac_splice import EDGE_FRAMES, PREROLL_FRAMES, splice_soundtrack, spliceable_aac_layout +from invokeai.app.util.mp4_metadata import find_box, write_mp4_tags +from invokeai.app.util.video_audio import extract_audio_pcm +from invokeai.app.util.video_encoding import make_mp4_writer, write_stereo_wav + +FPS = 16.0 +RATE = 32000 +SIZE = 64 +# How far into a trimmed soundtrack its encoded lead-in and its seam reach, and how far from its +# end the tail's: the copied run starts within a frame after EDGE_FRAMES of lead-in, and ends at +# least EDGE_FRAMES (and less than one more) before the end. +LEAD_IN_LIMIT = (EDGE_FRAMES + 2) * 1024 +TAIL_LIMIT = (EDGE_FRAMES + 1) * 1024 + + +def _signal(kind: str, n_frames: int, seed: int) -> np.ndarray: + n = round(n_frames / FPS * RATE) + t = np.arange(n) / RATE + if kind == "tone": + mono = 0.4 * np.sin(2 * np.pi * (440 + 110 * seed) * t) + else: + # Broadband noise is where the encoder would use perceptual noise substitution. + mono = 0.25 * np.random.default_rng(seed).standard_normal(n) + return np.stack([mono, 0.9 * mono]).astype(np.float32) + + +def _clip(dirpath: Path, name: str, n_frames: int, kind: str, seed: int) -> Path: + path, wav = dirpath / f"{name}.mp4", dirpath / f"{name}.wav" + write_stereo_wav(wav, _signal(kind, n_frames, seed), RATE) + writer = make_mp4_writer(path, FPS, audio_path=wav) + rng = np.random.default_rng(seed) + for _ in range(n_frames): + writer.append_data(rng.integers(0, 256, (SIZE, SIZE, 3), dtype=np.uint8)) + writer.close() + return path + + +def _save_to(saved: Path): + def _save(source_path: Path, **_: object) -> MagicMock: + # The stream-copy remux save() runs to embed the record. + write_mp4_tags(source_path, saved, {"invokeai_metadata": "{}"}) + return MagicMock() + + return _save + + +def _trim(dirpath: Path, source: Path, start: int, end: int, name: str = "trimmed") -> tuple[Path, MagicMock]: + saved = dirpath / f"{name}.mp4" + context = MagicMock() + context.videos.get_path.return_value = source + context.util.is_canceled.return_value = False + context.videos.save.side_effect = _save_to(saved) + node = ExtractVideoRangeInvocation(video=VideoField(video_name=source.stem), start_frame=start, end_frame=end) + base = MagicMock(video=VideoField(video_name=name), width=SIZE, height=SIZE, num_frames=1, fps=FPS, duration=1.0) + with patch("invokeai.app.invocations.video_frame_extract_range.VideoOutput.build", return_value=base): + node.invoke(context) + return saved, context + + +def _logged(context: MagicMock) -> str: + return "\n".join(str(call.args[0]) for call in context.logger.info.call_args_list) + + +def _movie_duration(path: Path) -> float: + with open(path, "rb") as fh: + fh.seek(0, 2) + moov = find_box(fh, 0, fh.tell(), b"moov") + assert moov is not None + mvhd = find_box(fh, *moov, b"mvhd") + assert mvhd is not None + fh.seek(mvhd[0]) + wide = fh.read(1)[0] == 1 + fh.seek(mvhd[0] + (20 if wide else 12)) + timescale, duration = struct.unpack(">IQ" if wide else ">II", fh.read(12 if wide else 8)) + return duration / timescale + + +def _window(source: Path, start: int, end: int) -> np.ndarray: + decoded = extract_audio_pcm(source) + assert decoded is not None + return decoded[0][:, round(start * RATE / FPS) : round((end + 1) * RATE / FPS)] + + +@pytest.mark.parametrize("kind", ["tone", "noise"]) +@pytest.mark.parametrize("start,end", [(20, 90), (0, 80), (33, 99)]) +def test_trim_copies_the_source_audio_it_keeps(tmp_path, monkeypatch, kind, start, end): + source = _clip(tmp_path, "source", 100, kind, seed=1) + trimmed, context = _trim(tmp_path, source, start, end) + assert "audio packets" in _logged(context) + + want = _window(source, start, end) + got = extract_audio_pcm(trimmed) + assert got is not None + n = end - start + 1 + # The soundtrack ends with the video, and decodes to whole frames past it. + assert _movie_duration(trimmed) == pytest.approx(n / FPS, abs=0.002) + assert 0 <= got[0].shape[1] - round(n * RATE / FPS) < 1024 + # Every frame away from the lead-in and tail seams is the source's own audio, bit for bit. + interior = slice(LEAD_IN_LIMIT, want.shape[1] - TAIL_LIMIT) + assert np.array_equal(got[0][:, interior], want[:, interior]) + + # The seams are no worse than re-encoding the whole range. + monkeypatch.setattr(ExtractVideoRangeInvocation, "_mux_spliced_audio", lambda *args, **kwargs: False) + full, _ = _trim(tmp_path, source, start, end, name="full") + reference = extract_audio_pcm(full) + assert reference is not None + m = min(got[0].shape[1], reference[0].shape[1], want.shape[1]) + spliced_error = np.abs(got[0][:, :m] - want[:, :m]) + full_error = np.abs(reference[0][:, :m] - want[:, :m]) + for seam in (slice(0, LEAD_IN_LIMIT), slice(want.shape[1] - TAIL_LIMIT, m)): + assert ( + float(np.sqrt(np.mean(spliced_error[:, seam] ** 2))) + <= 2 * float(np.sqrt(np.mean(full_error[:, seam] ** 2))) + 1e-3 + ) + + +def test_copied_audio_that_does_not_decode_back_is_rejected(tmp_path): + """The splice trusts nothing it has not decoded: if the source's packets do not reproduce + the soundtrack's samples where they are copied (here, a timeline nudged by one LSB), the + splice is refused and the caller re-encodes.""" + source = _clip(tmp_path, "source", 100, "tone", seed=1) + video = tmp_path / "video.mp4" + writer = make_mp4_writer(video, FPS) + for _ in range(71): + writer.append_data(np.zeros((SIZE, SIZE, 3), dtype=np.uint8)) + writer.close() + layout = spliceable_aac_layout(source, RATE) + assert layout is not None + timeline = _window(source, 20, 90) + window_start = round(20 * RATE / FPS) + first = -(-(window_start + layout.priming + EDGE_FRAMES * 1024) // 1024) + copied_at = first * 1024 - layout.priming - window_start + end = (round(91 * RATE / FPS) - EDGE_FRAMES * 1024 + layout.priming) // 1024 + args = (video, source, layout.priming, first, end, copied_at) + + assert splice_soundtrack(*args, timeline, RATE, tmp_path / "ok.mp4") is None + nudged = timeline.copy() + nudged[:, timeline.shape[1] // 2] += 1 / 32768 + assert ( + splice_soundtrack(*args, nudged, RATE, tmp_path / "bad.mp4") == "the copied audio did not decode back exactly" + ) + + +def test_extending_a_trimmed_source_keeps_its_audio(tmp_path): + """The Extend graph: trim the source, then join a continuation onto the trim.""" + source = _clip(tmp_path, "source", 100, "tone", seed=1) + trimmed, trim_context = _trim(tmp_path, source, 20, 95) + assert "audio packets" in _logged(trim_context) + continuation = _clip(tmp_path, "continuation", 40, "tone", seed=3) + + joined = tmp_path / "joined.mp4" + by_name = {"trimmed": trimmed, "continuation": continuation} + context = MagicMock() + context.videos.get_path.side_effect = lambda name: by_name[name] + context.util.is_canceled.return_value = False + context.videos.save.side_effect = _save_to(joined) + node = VideoConcatInvocation( + videos=[VideoField(video_name="trimmed"), VideoField(video_name="continuation")], + transition="crossfade", + transition_frames=2, + ) + with patch("invokeai.app.invocations.video_concat.VideoOutput.build", return_value=MagicMock()): + node.invoke(context) + assert "audio samples of the first clip without re-encoding" in _logged(context) + + want = _window(source, 20, 95) + got = extract_audio_pcm(joined) + assert got is not None + # The source's audio inside the trim, clear of the trim's lead-in and the join, is unchanged. + interior = slice(LEAD_IN_LIMIT, want.shape[1] - round(2 * RATE / FPS) - (PREROLL_FRAMES + 2) * 1024) + assert np.array_equal(got[0][:, interior], want[:, interior]) + reader = iio2.get_reader(str(joined), format="FFMPEG") + try: + assert sum(1 for _ in reader) == 76 + 40 - 2 + finally: + reader.close() + + +def test_trim_splices_when_frames_hold_fractional_samples(tmp_path): + """At 44.1 kHz and 24 fps a frame is 1837.5 samples: rounding the trimmed span on its own + came out one off the window, which interpolated the whole soundtrack and ruled out a copy.""" + rate, fps = 44100, 24.0 + n = round(110 / fps * rate) + tone = (0.4 * np.sin(2 * np.pi * 523.0 * np.arange(n) / rate)).astype(np.float32) + wav, source = tmp_path / "source.wav", tmp_path / "source.mp4" + write_stereo_wav(wav, np.stack([tone, tone]), rate) + writer = make_mp4_writer(source, fps, audio_path=wav) + for i in range(110): + writer.append_data(np.full((SIZE, SIZE, 3), i, dtype=np.uint8)) + writer.close() + + trimmed, context = _trim(tmp_path, source, 7, 100) + assert "audio packets" in _logged(context) + decoded, got = extract_audio_pcm(source), extract_audio_pcm(trimmed) + assert decoded is not None and got is not None + want = decoded[0][:, round(7 * rate / fps) : round(101 * rate / fps)] + interior = slice(LEAD_IN_LIMIT, want.shape[1] - TAIL_LIMIT) + assert np.array_equal(got[0][:, interior], want[:, interior]) From ca9837f4f4a2c4ffa734bb50add95e394dbedafe Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Tue, 13 Oct 2026 16:40:00 -0400 Subject: [PATCH 4/7] fix(video): bound splice reads and walks, cancel audio decodes Sample entries are size-capped before any read and every container walk is box-capped, so crafted uploads cannot exhaust memory or hold a worker. One cancellable ffmpeg runner now backs the splice copies, audio decodes and muxes, and the passthrough tag remux. --- invokeai/app/invocations/video_concat.py | 4 +- .../invocations/video_frame_extract_range.py | 6 +- invokeai/app/util/aac_splice.py | 63 ++++++--- invokeai/app/util/h264_splice.py | 105 +++++++-------- invokeai/app/util/mp4_metadata.py | 25 ++-- invokeai/app/util/video_audio.py | 90 +++++++++---- .../invocations/test_video_concat_splice.py | 18 ++- tests/app/util/test_h264_splice.py | 123 ++++++++++++++++-- 8 files changed, 294 insertions(+), 140 deletions(-) diff --git a/invokeai/app/invocations/video_concat.py b/invokeai/app/invocations/video_concat.py index cb4dfe5126b..0d7a48957aa 100644 --- a/invokeai/app/invocations/video_concat.py +++ b/invokeai/app/invocations/video_concat.py @@ -281,7 +281,7 @@ def _save_joined( ): wav_path = work_dir / "audio.wav" write_stereo_wav(wav_path, pcm, rate) - mux_audio_into_video(tmp_path, wav_path, muxed_path) + mux_audio_into_video(tmp_path, wav_path, muxed_path, context.util.is_canceled) source_path = muxed_path duration = num_frames / output_fps @@ -557,7 +557,7 @@ def _build_audio_track( for path in paths: if context.util.is_canceled(): raise CanceledException - extracted.append(extract_audio_pcm(path)) + extracted.append(extract_audio_pcm(path, is_canceled=context.util.is_canceled)) if all(item is None for item in extracted): return None if context.util.is_canceled(): diff --git a/invokeai/app/invocations/video_frame_extract_range.py b/invokeai/app/invocations/video_frame_extract_range.py index b71a9e919b8..c14a8189ee8 100644 --- a/invokeai/app/invocations/video_frame_extract_range.py +++ b/invokeai/app/invocations/video_frame_extract_range.py @@ -299,7 +299,7 @@ def invoke(self, context: InvocationContext) -> ExtractVideoRangeOutput: wav_tmp.close() wav_path = Path(wav_tmp.name) write_stereo_wav(wav_path, pcm, rate) - mux_audio_into_video(tmp_path, wav_path, muxed_path) + mux_audio_into_video(tmp_path, wav_path, muxed_path, context.util.is_canceled) source_path = muxed_path out_duration = num_frames / output_fps @@ -452,7 +452,7 @@ def _save_unchanged( copy_path = Path(work) / "source.mp4" # A stream-copy remux that drops the source's keyed metadata: save() writes this # node's record, and if that rewrite fails the source's must not be read back as it. - write_mp4_tags(video_path, copy_path, {}) + write_mp4_tags(video_path, copy_path, {}, context.util.is_canceled) # A cancellation that arrived during an uninterruptible stage must not be saved. if context.util.is_canceled(): raise CanceledException @@ -564,7 +564,7 @@ def _build_audio_track( (matching what the retime does to the video). """ context.util.signal_progress("Extracting audio") - extracted = extract_audio_pcm(video_path) + extracted = extract_audio_pcm(video_path, is_canceled=context.util.is_canceled) if extracted is None: return None if context.util.is_canceled(): diff --git a/invokeai/app/util/aac_splice.py b/invokeai/app/util/aac_splice.py index e31d140a029..c46216293a1 100644 --- a/invokeai/app/util/aac_splice.py +++ b/invokeai/app/util/aac_splice.py @@ -47,7 +47,13 @@ # start (a lead-in) or its end-of-stream flush (a tail). EDGE_FRAMES = PREROLL_FRAMES + 1 _MAX_TOP_LEVEL_BOXES = 1024 +# Children walked per container (moov, trak, mdia, minf, stbl, edts): real ones hold about a +# dozen, and the boxes are the uploader's, so a flood is given up on rather than scanned. +_MAX_CHILD_BOXES = 64 _MAX_PACKETS = 1 << 24 +# A real AAC sample entry (esds with a two-to-five-byte AudioSpecificConfig, btrt) is well +# under a hundred bytes. +_MAX_SAMPLE_ENTRY_BYTES = 4 * 1024 # AudioSampleEntry fields before its child boxes (SampleEntry 8 + version-0 audio header 20). _AUDIO_SAMPLE_ENTRY_HEADER = 28 _ES_DESCRIPTOR, _DECODER_CONFIG_DESCRIPTOR, _DECODER_SPECIFIC_INFO = 3, 4, 5 @@ -122,11 +128,11 @@ def read_aac_layout(path: Path) -> Optional[AacLayout]: moov = find_box(fh, 0, fh.tell(), b"moov", max_boxes=_MAX_TOP_LEVEL_BOXES) if moov is None: return None - for box_type, trak_start, trak_end in iter_boxes(fh, *moov): + for box_type, trak_start, trak_end in iter_boxes(fh, *moov, max_boxes=_MAX_CHILD_BOXES): if box_type != b"trak": continue - mdia = find_box(fh, trak_start, trak_end, b"mdia") - hdlr = find_box(fh, *mdia, b"hdlr") if mdia is not None else None + mdia = find_box(fh, trak_start, trak_end, b"mdia", max_boxes=_MAX_CHILD_BOXES) + hdlr = find_box(fh, *mdia, b"hdlr", max_boxes=_MAX_CHILD_BOXES) if mdia is not None else None if hdlr is None or hdlr[0] + 12 > hdlr[1]: continue fh.seek(hdlr[0] + 8) @@ -137,16 +143,16 @@ def read_aac_layout(path: Path) -> Optional[AacLayout]: def _read_audio_track(fh: BinaryIO, trak: tuple[int, int], mdia: tuple[int, int]) -> Optional[AacLayout]: - mdhd = find_box(fh, *mdia, b"mdhd") - minf = find_box(fh, *mdia, b"minf") - stbl = find_box(fh, *minf, b"stbl") if minf is not None else None + mdhd = find_box(fh, *mdia, b"mdhd", max_boxes=_MAX_CHILD_BOXES) + minf = find_box(fh, *mdia, b"minf", max_boxes=_MAX_CHILD_BOXES) + stbl = find_box(fh, *minf, b"stbl", max_boxes=_MAX_CHILD_BOXES) if minf is not None else None if mdhd is None or stbl is None or mdhd[0] + 24 > mdhd[1]: return None fh.seek(mdhd[0]) fh.seek(mdhd[0] + (20 if fh.read(1)[0] == 1 else 12)) (sample_rate,) = struct.unpack(">I", fh.read(4)) - stsd = find_box(fh, *stbl, b"stsd") + stsd = find_box(fh, *stbl, b"stsd", max_boxes=_MAX_CHILD_BOXES) if stsd is None or stsd[0] + 8 > stsd[1]: return None fh.seek(stsd[0] + 4) @@ -156,6 +162,10 @@ def _read_audio_track(fh: BinaryIO, trak: tuple[int, int], mdia: tuple[int, int] if not entries or entries[0][0] != b"mp4a": return None _, entry_start, entry_end = entries[0] + # Everything read below — the esds above all — lies inside the entry, which is the + # uploader's: bounded before any of it is read. + if entry_end - entry_start > _MAX_SAMPLE_ENTRY_BYTES: + return None fh.seek(entry_start + 8) if fh.read(2) != b"\x00\x00": # QuickTime v1/v2 sound entries lay their fields out differently return None @@ -179,7 +189,7 @@ def _read_audio_track(fh: BinaryIO, trak: tuple[int, int], mdia: tuple[int, int] if len(decoder_config) == 5 and decoder_config.endswith(_NO_SBR_EXTENSION): decoder_config = decoder_config[:2] - stts = find_box(fh, *stbl, b"stts") + stts = find_box(fh, *stbl, b"stts", max_boxes=_MAX_CHILD_BOXES) if stts is None or stts[0] + 8 > stts[1]: return None fh.seek(stts[0] + 4) @@ -195,8 +205,8 @@ def _read_audio_track(fh: BinaryIO, trak: tuple[int, int], mdia: tuple[int, int] if packet_count > _MAX_PACKETS: return None - edts = find_box(fh, *trak, b"edts") - elst = find_box(fh, *edts, b"elst") if edts is not None else None + edts = find_box(fh, *trak, b"edts", max_boxes=_MAX_CHILD_BOXES) + elst = find_box(fh, *edts, b"elst", max_boxes=_MAX_CHILD_BOXES) if edts is not None else None priming = 0 if elst is not None: if elst[0] + 8 > elst[1]: @@ -275,7 +285,7 @@ def _set_audio_edit(path: Path, media_time: int, samples: int, sample_rate: int) with open(path, "r+b") as fh: fh.seek(0, 2) moov = find_box(fh, 0, fh.tell(), b"moov", max_boxes=_MAX_TOP_LEVEL_BOXES) - mvhd = find_box(fh, *moov, b"mvhd") if moov is not None else None + mvhd = find_box(fh, *moov, b"mvhd", max_boxes=_MAX_CHILD_BOXES) if moov is not None else None if mvhd is None: raise SpliceError(f"{path.name} has no movie header") fh.seek(mvhd[0]) @@ -285,12 +295,12 @@ def _set_audio_edit(path: Path, media_time: int, samples: int, sample_rate: int) duration = round(samples * movie_timescale / sample_rate) track_durations: list[int] = [] patched = False - for box_type, trak_start, trak_end in iter_boxes(fh, *moov): + for box_type, trak_start, trak_end in iter_boxes(fh, *moov, max_boxes=_MAX_CHILD_BOXES): if box_type != b"trak": continue - tkhd = find_box(fh, trak_start, trak_end, b"tkhd") - mdia = find_box(fh, trak_start, trak_end, b"mdia") - hdlr = find_box(fh, *mdia, b"hdlr") if mdia is not None else None + tkhd = find_box(fh, trak_start, trak_end, b"tkhd", max_boxes=_MAX_CHILD_BOXES) + mdia = find_box(fh, trak_start, trak_end, b"mdia", max_boxes=_MAX_CHILD_BOXES) + hdlr = find_box(fh, *mdia, b"hdlr", max_boxes=_MAX_CHILD_BOXES) if mdia is not None else None if tkhd is None or hdlr is None: raise SpliceError(f"{path.name} has a malformed track") fh.seek(hdlr[0] + 8) @@ -299,8 +309,8 @@ def _set_audio_edit(path: Path, media_time: int, samples: int, sample_rate: int) wide = fh.read(1)[0] == 1 tkhd_duration_at, field = tkhd[0] + (28 if wide else 20), ">Q" if wide else ">I" if is_target and not patched: - edts = find_box(fh, trak_start, trak_end, b"edts") - elst = find_box(fh, *edts, b"elst") if edts is not None else None + edts = find_box(fh, trak_start, trak_end, b"edts", max_boxes=_MAX_CHILD_BOXES) + elst = find_box(fh, *edts, b"elst", max_boxes=_MAX_CHILD_BOXES) if edts is not None else None if elst is None: raise SpliceError(f"{path.name}'s {handler.decode()} track has no edit list") fh.seek(elst[0]) @@ -328,7 +338,7 @@ def _encode( wav, adts = work / f"{name}.wav", work / f"{name}.aac" write_stereo_wav(wav, pcm, sample_rate) run_ffmpeg(["-i", str(wav), *AAC_ENCODE_ARGS, "-f", "adts", str(adts)], "encoding the new audio", is_canceled) - decoded = extract_audio_pcm(adts) + decoded = extract_audio_pcm(adts, is_canceled=is_canceled) if decoded is None: raise SpliceError("the new audio decoded to nothing") return _adts_frames(adts.read_bytes()), decoded[0] @@ -398,7 +408,20 @@ def splice_soundtrack( # The source's packets as ADTS, sliced by index: an ADTS frame per packet. source_adts = work_dir / "source.aac" run_ffmpeg( - ["-i", str(source), "-map", "0:a:0", "-c", "copy", "-f", "adts", str(source_adts)], + # Only the packets up to the copied run's end are needed. + [ + "-i", + str(source), + "-map", + "0:a:0", + "-c", + "copy", + "-frames:a", + str(end), + "-f", + "adts", + str(source_adts), + ], f"copying the audio packets of {source.name}", is_canceled, ) @@ -463,7 +486,7 @@ def splice_soundtrack( is_canceled, ) _set_audio_edit(out_path, media_time, timeline.shape[1], sample_rate) - muxed = extract_audio_pcm(out_path) + muxed = extract_audio_pcm(out_path, is_canceled=is_canceled) except (SpliceError, AudioExtractionError, OSError) as e: return f"it cannot be stream-copied ({e})" diff --git a/invokeai/app/util/h264_splice.py b/invokeai/app/util/h264_splice.py index 33712d45dca..af35c267a2e 100644 --- a/invokeai/app/util/h264_splice.py +++ b/invokeai/app/util/h264_splice.py @@ -38,30 +38,31 @@ import struct import subprocess import tempfile -import time from pathlib import Path from typing import BinaryIO, Callable, NamedTuple, Optional -import imageio_ffmpeg import numpy as np -from invokeai.app.services.session_processor.session_processor_common import CanceledException from invokeai.app.util.mp4_metadata import find_box, iter_boxes +from invokeai.app.util.video_audio import run_ffmpeg_process # A stream copy is disk-bound; this only stops a wedged child from holding the node forever. SPLICE_TIMEOUT_SECONDS = 600 -# How often a running ffmpeg child is checked for cancellation. -_CANCEL_POLL_SECONDS = 0.1 # ``moov`` precedes or follows ``mdat`` in a real file; a dozen top-level boxes at most. _MAX_TOP_LEVEL_BOXES = 1024 +# Children walked per container (moov, trak, mdia, minf, stbl, edts): real ones hold about a +# dozen, and the boxes are the uploader's, so a flood is given up on rather than scanned. +_MAX_CHILD_BOXES = 64 # Keyframes examined per track: at the 2 s interval this app writes, nine hours of video. _MAX_SYNC_ENTRIES = 1 << 14 # VisualSampleEntry fields before its child boxes (SampleEntry 8 + visual header 70), and where # among them the 32-byte compressor name sits. _VISUAL_SAMPLE_ENTRY_HEADER = 78 _COMPRESSOR_NAME = slice(42, 74) -# A real sample entry holds a handful of boxes (avcC, btrt, colr, pasp...); more is hostile. +# A real sample entry holds a handful of boxes (avcC, btrt, colr, pasp...) in a few hundred +# bytes; more of either is hostile, and refused before any of it is read. _MAX_ENTRY_CHILDREN = 32 +_MAX_SAMPLE_ENTRY_BYTES = 64 * 1024 # NAL units examined per keyframe before its first slice (AUD, SEI, SPS, PPS precede it). _MAX_NALS_BEFORE_SLICE = 64 # Shortest part x264 (bframes=3, as make_mp4_writer leaves it) encodes with its full B-frame @@ -121,30 +122,30 @@ def read_h264_layout(path: Path) -> Optional[H264Layout]: moov = find_box(fh, 0, file_size, b"moov", max_boxes=_MAX_TOP_LEVEL_BOXES) if moov is None: return None - for box_type, trak_start, trak_end in iter_boxes(fh, *moov): + for box_type, trak_start, trak_end in iter_boxes(fh, *moov, max_boxes=_MAX_CHILD_BOXES): if box_type != b"trak": continue - mdia = find_box(fh, trak_start, trak_end, b"mdia") + mdia = find_box(fh, trak_start, trak_end, b"mdia", max_boxes=_MAX_CHILD_BOXES) if mdia is None: continue - hdlr = find_box(fh, *mdia, b"hdlr") + hdlr = find_box(fh, *mdia, b"hdlr", max_boxes=_MAX_CHILD_BOXES) if hdlr is None or hdlr[0] + 12 > hdlr[1]: continue fh.seek(hdlr[0] + 8) if fh.read(4) != b"vide": continue # First video track only: that is the stream every ffmpeg call here maps (0:v:0). - minf = find_box(fh, *mdia, b"minf") - stbl = find_box(fh, *minf, b"stbl") if minf is not None else None - movie_timescale = _read_timescale(fh, find_box(fh, *moov, b"mvhd")) - media_timescale = _read_timescale(fh, find_box(fh, *mdia, b"mdhd")) + minf = find_box(fh, *mdia, b"minf", max_boxes=_MAX_CHILD_BOXES) + stbl = find_box(fh, *minf, b"stbl", max_boxes=_MAX_CHILD_BOXES) if minf is not None else None + movie_timescale = _read_timescale(fh, find_box(fh, *moov, b"mvhd", max_boxes=_MAX_CHILD_BOXES)) + media_timescale = _read_timescale(fh, find_box(fh, *mdia, b"mdhd", max_boxes=_MAX_CHILD_BOXES)) if stbl is None or movie_timescale is None or media_timescale is None: return None - if not _has_identity_matrix(fh, find_box(fh, trak_start, trak_end, b"tkhd")): + if not _has_identity_matrix(fh, find_box(fh, trak_start, trak_end, b"tkhd", max_boxes=_MAX_CHILD_BOXES)): return None - edts = find_box(fh, trak_start, trak_end, b"edts") + edts = find_box(fh, trak_start, trak_end, b"edts", max_boxes=_MAX_CHILD_BOXES) if edts is not None and not _edits_show_every_sample( - fh, find_box(fh, *edts, b"elst"), stbl, movie_timescale, media_timescale + fh, find_box(fh, *edts, b"elst", max_boxes=_MAX_CHILD_BOXES), stbl, movie_timescale, media_timescale ): return None return _read_video_stbl(fh, *stbl) @@ -208,7 +209,7 @@ def _edits_show_every_sample( first_delta = runs[0][1] first_offset = 0 - ctts = find_box(fh, *stbl, b"ctts") + ctts = find_box(fh, *stbl, b"ctts", max_boxes=_MAX_CHILD_BOXES) if ctts is not None: if ctts[0] + 16 > ctts[1]: return False @@ -223,7 +224,7 @@ def _edits_show_every_sample( def _read_stts(fh: BinaryIO, stbl: tuple[int, int]) -> Optional[list[tuple[int, int]]]: """The ``(sample_count, sample_delta)`` runs of the time-to-sample table.""" - stts = find_box(fh, *stbl, b"stts") + stts = find_box(fh, *stbl, b"stts", max_boxes=_MAX_CHILD_BOXES) if stts is None or stts[0] + 8 > stts[1]: return None fh.seek(stts[0] + 4) @@ -260,8 +261,11 @@ def _sample_locations( fh: BinaryIO, stbl: tuple[int, int], sizes: np.ndarray, wanted: tuple[int, ...], file_size: int ) -> Optional[list[tuple[int, int]]]: """``(file_offset, size)`` of each sample in ``wanted`` (ascending), from the chunk tables.""" - stsc = find_box(fh, *stbl, b"stsc") - stco, co64 = find_box(fh, *stbl, b"stco"), find_box(fh, *stbl, b"co64") + stsc = find_box(fh, *stbl, b"stsc", max_boxes=_MAX_CHILD_BOXES) + stco, co64 = ( + find_box(fh, *stbl, b"stco", max_boxes=_MAX_CHILD_BOXES), + find_box(fh, *stbl, b"co64", max_boxes=_MAX_CHILD_BOXES), + ) offsets_box, offset_format = (stco, "I") if stco is not None else (co64, "Q") # Each table opens with version/flags and a 4-byte count. if stsc is None or offsets_box is None or any(box[0] + 8 > box[1] for box in (stsc, offsets_box)): @@ -326,8 +330,8 @@ def _is_idr(fh: BinaryIO, offset: int, size: int, length_size: int) -> bool: def _read_video_stbl(fh: BinaryIO, stbl_start: int, stbl_end: int) -> Optional[H264Layout]: stbl = (stbl_start, stbl_end) - stsd = find_box(fh, *stbl, b"stsd") - stsz = find_box(fh, *stbl, b"stsz") + stsd = find_box(fh, *stbl, b"stsd", max_boxes=_MAX_CHILD_BOXES) + stsz = find_box(fh, *stbl, b"stsz", max_boxes=_MAX_CHILD_BOXES) if stsd is None or stsz is None: return None @@ -340,7 +344,15 @@ def _read_video_stbl(fh: BinaryIO, stbl_start: int, stbl_end: int) -> Optional[H if not entries or entries[0][0] != b"avc1": return None _, entry_start, entry_end = entries[0] - avcc = find_box(fh, entry_start + _VISUAL_SAMPLE_ENTRY_HEADER, entry_end, b"avcC") + # Everything read below lies inside the entry, so one size bound covers every child box. + if entry_end - entry_start > _MAX_SAMPLE_ENTRY_BYTES: + return None + children = list( + iter_boxes(fh, entry_start + _VISUAL_SAMPLE_ENTRY_HEADER, entry_end, max_boxes=_MAX_ENTRY_CHILDREN + 1) + ) + if len(children) > _MAX_ENTRY_CHILDREN: + return None + avcc = next(((start, end) for child, start, end in children if child == b"avcC"), None) if avcc is None or avcc[0] >= avcc[1]: return None fh.seek(avcc[0]) @@ -351,11 +363,6 @@ def _read_video_stbl(fh: BinaryIO, stbl_start: int, stbl_end: int) -> Optional[H # an ffmpeg upgrade still match what it writes after one. header_fields[_COMPRESSOR_NAME] = bytes(_COMPRESSOR_NAME.stop - _COMPRESSOR_NAME.start) parts = [bytes(header_fields)] - children = list( - iter_boxes(fh, entry_start + _VISUAL_SAMPLE_ENTRY_HEADER, entry_end, max_boxes=_MAX_ENTRY_CHILDREN + 1) - ) - if len(children) > _MAX_ENTRY_CHILDREN: - return None for child, child_start, child_end in children: if child != b"btrt": fh.seek(child_start) @@ -372,7 +379,7 @@ def _read_video_stbl(fh: BinaryIO, stbl_start: int, stbl_end: int) -> Optional[H return None sample_count = len(sizes) - stss = find_box(fh, *stbl, b"stss") + stss = find_box(fh, *stbl, b"stss", max_boxes=_MAX_CHILD_BOXES) if stss is None: # No sync-sample table means every sample is a sync sample (an all-intra stream), each # of which would be examined below: only a short one is worth splicing. @@ -416,38 +423,16 @@ def _read_video_stbl(fh: BinaryIO, stbl_start: int, stbl_end: int) -> Optional[H def run_ffmpeg(args: list[str], description: str, is_canceled: Optional[Callable[[], bool]] = None) -> None: """Run the bundled ffmpeg, raising SpliceError if it fails or outlasts the timeout. - ``is_canceled`` is polled while the child runs; a cancellation kills and reaps it and - raises CanceledException, so a canceled node never waits on a long copy. + ``is_canceled`` stops the child promptly and raises CanceledException (see + ``run_ffmpeg_process``), so a canceled node never waits on a long copy. """ - if is_canceled is not None and is_canceled(): - raise CanceledException - with tempfile.TemporaryFile() as stderr: - proc = subprocess.Popen( - [imageio_ffmpeg.get_ffmpeg_exe(), "-y", "-loglevel", "error", *args], - stdin=subprocess.DEVNULL, - stdout=subprocess.DEVNULL, - stderr=stderr, - ) - deadline = time.monotonic() + SPLICE_TIMEOUT_SECONDS - try: - while True: - try: - proc.wait(timeout=_CANCEL_POLL_SECONDS) - break - except subprocess.TimeoutExpired: - pass - if is_canceled is not None and is_canceled(): - raise CanceledException - if time.monotonic() > deadline: - raise SpliceError(f"ffmpeg timed out {description}") - finally: - if proc.returncode is None: - proc.kill() - proc.wait() - if proc.returncode != 0: - stderr.seek(0) - message = stderr.read().decode("utf-8", errors="replace").strip() - raise SpliceError(f"ffmpeg failed {description}: {message[-500:]}") + try: + proc = run_ffmpeg_process(args, SPLICE_TIMEOUT_SECONDS, is_canceled) + except subprocess.TimeoutExpired as e: + raise SpliceError(f"ffmpeg timed out {description}") from e + if proc.returncode != 0: + message = proc.stderr.decode("utf-8", errors="replace").strip() + raise SpliceError(f"ffmpeg failed {description}: {message[-500:]}") def copy_frames( diff --git a/invokeai/app/util/mp4_metadata.py b/invokeai/app/util/mp4_metadata.py index d4b1a9759d9..1074ff6b75c 100644 --- a/invokeai/app/util/mp4_metadata.py +++ b/invokeai/app/util/mp4_metadata.py @@ -25,9 +25,7 @@ import tempfile from collections.abc import Collection, Iterator, Mapping from pathlib import Path -from typing import BinaryIO, Optional - -import imageio_ffmpeg +from typing import BinaryIO, Callable, Optional # A single tag is a JSON string the app itself produced; a graph JSON is the largest and # runs to a few hundred KiB. The cap exists only so a hostile file cannot make the reader @@ -52,10 +50,6 @@ class Mp4MetadataError(RuntimeError): """The MP4 could not be rewritten with the requested tags.""" -def _ffmpeg_exe() -> str: - return imageio_ffmpeg.get_ffmpeg_exe() - - def iter_boxes(fh: BinaryIO, start: int, end: int, max_boxes: Optional[int] = None) -> Iterator[tuple[bytes, int, int]]: """Yield ``(type, payload_start, payload_end)`` for each box in ``[start, end)``. @@ -220,7 +214,9 @@ def _escape_ffmetadata(value: str) -> str: ) -def write_mp4_tags(src: Path, dst: Path, tags: Mapping[str, str]) -> None: +def write_mp4_tags( + src: Path, dst: Path, tags: Mapping[str, str], is_canceled: Optional[Callable[[], bool]] = None +) -> None: """Write ``dst`` = ``src``'s first video stream and every audio stream, stream-copied, with ``tags`` as the file's keyed metadata (replacing any global tags ``src`` had). An empty mapping therefore strips ``src``'s global tags. @@ -243,12 +239,11 @@ def write_mp4_tags(src: Path, dst: Path, tags: Mapping[str, str]) -> None: ffmeta_path = Path(ffmeta.name) try: try: - proc = subprocess.run( + # Imported here: video_audio is a heavier module than this one's other callers need. + from invokeai.app.util.video_audio import run_ffmpeg_process + + proc = run_ffmpeg_process( [ - _ffmpeg_exe(), - "-y", - "-loglevel", - "error", "-i", str(src), "-i", @@ -267,8 +262,8 @@ def write_mp4_tags(src: Path, dst: Path, tags: Mapping[str, str]) -> None: "mp4", str(dst), ], - capture_output=True, - timeout=REMUX_TIMEOUT_SECONDS, + REMUX_TIMEOUT_SECONDS, + is_canceled, ) except subprocess.TimeoutExpired as e: raise Mp4MetadataError(f"ffmpeg timed out writing metadata into {dst.name}") from e diff --git a/invokeai/app/util/video_audio.py b/invokeai/app/util/video_audio.py index eca1db28281..2a28b259d10 100644 --- a/invokeai/app/util/video_audio.py +++ b/invokeai/app/util/video_audio.py @@ -18,11 +18,19 @@ import subprocess import tempfile +import time import wave from pathlib import Path +from typing import Callable, Optional import numpy as np +from invokeai.app.services.session_processor.session_processor_common import CanceledException + +# How often a running ffmpeg child is checked for cancellation, and how much of its stderr is kept. +_CANCEL_POLL_SECONDS = 0.1 +_STDERR_TAIL_BYTES = 8 * 1024 + # How every soundtrack the app writes is encoded; ``aac_splice`` matches a spliced tail to it. # Perceptual noise substitution is off: the decoder fills those bands from a random generator # whose state carries across packets, so a freshly encoded lead-in that used it would change @@ -37,11 +45,50 @@ def _ffmpeg_exe() -> str: return imageio_ffmpeg.get_ffmpeg_exe() +def run_ffmpeg_process( + args: list[str], timeout: float, is_canceled: Optional[Callable[[], bool]] = None +) -> subprocess.CompletedProcess[bytes]: + """Run the bundled ffmpeg (quiet, overwriting) with ``args``; returns its exit status and stderr. + + The child is polled rather than waited on, so ``is_canceled`` can stop it: a cancellation + (checked before it starts, too) kills and reaps it and raises CanceledException, and + outliving ``timeout`` kills and reaps it and raises ``subprocess.TimeoutExpired``. A + canceled node therefore never waits on a long decode, encode or copy. + """ + if is_canceled is not None and is_canceled(): + raise CanceledException + command = [_ffmpeg_exe(), "-y", "-loglevel", "error", *args] + with tempfile.TemporaryFile() as stderr: + proc = subprocess.Popen(command, stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL, stderr=stderr) + deadline = time.monotonic() + timeout + try: + while True: + try: + proc.wait(timeout=_CANCEL_POLL_SECONDS) + break + except subprocess.TimeoutExpired: + pass + if is_canceled is not None and is_canceled(): + raise CanceledException + if time.monotonic() > deadline: + raise subprocess.TimeoutExpired(command, timeout) + finally: + if proc.returncode is None: + proc.kill() + proc.wait() + # Only the tail is ever reported; a decoder logging per packet could write a great deal. + stderr.seek(0, 2) + stderr.seek(max(0, stderr.tell() - _STDERR_TAIL_BYTES)) + return subprocess.CompletedProcess(command, proc.returncode, b"", stderr.read()) + + class AudioExtractionError(RuntimeError): """ffmpeg failed to decode an audio track for a reason other than 'no audio stream'.""" -def extract_audio_pcm(video_path: Path, *, float_pcm: bool = False) -> tuple[np.ndarray, int] | None: +def extract_audio_pcm( + video_path: Path, *, float_pcm: bool = False, is_canceled: Optional[Callable[[], bool]] = None +) -> tuple[np.ndarray, int] | None: """Decode ``video_path``'s first audio stream to float32 stereo PCM at its native sample rate. Returns ``(samples, sample_rate)`` with ``samples`` shaped ``(2, n)`` in [-1, 1], or @@ -52,6 +99,8 @@ def extract_audio_pcm(video_path: Path, *, float_pcm: bool = False) -> tuple[np. With ``float_pcm=True`` the decode goes through ``pcm_f32le`` instead of ``pcm_s16le``, preserving the container's float samples without an int16 quantization. MiniMax H3 reference conditioning requires this: upstream conditions on the decoder's float output. + + ``is_canceled`` stops the decode promptly (see :func:`run_ffmpeg_process`). """ suffix = ".f32" if float_pcm else ".wav" with tempfile.NamedTemporaryFile(prefix="invokeai_audio_extract_", suffix=suffix, delete=False) as tmp: @@ -63,24 +112,10 @@ def extract_audio_pcm(video_path: Path, *, float_pcm: bool = False) -> tuple[np. ) try: try: - proc = subprocess.run( - [ - _ffmpeg_exe(), - "-y", - "-loglevel", - "error", - "-i", - str(video_path), - "-map", - "0:a:0?", - "-vn", - "-ac", - "2", - *codec_args, - str(out_path), - ], - capture_output=True, + proc = run_ffmpeg_process( + ["-i", str(video_path), "-map", "0:a:0?", "-vn", "-ac", "2", *codec_args, str(out_path)], timeout=600, + is_canceled=is_canceled, ) except subprocess.TimeoutExpired as e: raise AudioExtractionError(f"ffmpeg timed out extracting audio from {video_path.name}") from e @@ -164,13 +199,16 @@ def resample_linear(samples: np.ndarray, num_output_samples: int) -> np.ndarray: return np.stack([np.interp(positions, grid, samples[c]) for c in range(2)]).astype(np.float32) -def mux_audio_into_video(video_path: Path, wav_path: Path, output_path: Path) -> None: +def mux_audio_into_video( + video_path: Path, wav_path: Path, output_path: Path, is_canceled: Optional[Callable[[], bool]] = None +) -> None: """Produce ``output_path`` = ``video_path``'s streams with ``wav_path`` as an AAC track. The video stream is copied bit-exactly (no re-encode); only the audio is encoded. + ``is_canceled`` stops the mux promptly (see :func:`run_ffmpeg_process`). """ try: - proc = _run_mux(video_path, wav_path, output_path) + proc = _run_mux(video_path, wav_path, output_path, is_canceled) except subprocess.TimeoutExpired as e: raise AudioExtractionError(f"ffmpeg timed out muxing audio into {output_path.name}") from e if proc.returncode != 0: @@ -178,13 +216,11 @@ def mux_audio_into_video(video_path: Path, wav_path: Path, output_path: Path) -> raise AudioExtractionError(f"ffmpeg could not mux audio into {output_path.name}: {stderr.strip()[-500:]}") -def _run_mux(video_path: Path, wav_path: Path, output_path: Path) -> subprocess.CompletedProcess[bytes]: - return subprocess.run( +def _run_mux( + video_path: Path, wav_path: Path, output_path: Path, is_canceled: Optional[Callable[[], bool]] +) -> subprocess.CompletedProcess[bytes]: + return run_ffmpeg_process( [ - _ffmpeg_exe(), - "-y", - "-loglevel", - "error", "-i", str(video_path), "-i", @@ -198,6 +234,6 @@ def _run_mux(video_path: Path, wav_path: Path, output_path: Path) -> subprocess. *AAC_ENCODE_ARGS, str(output_path), ], - capture_output=True, timeout=600, + is_canceled=is_canceled, ) diff --git a/tests/app/invocations/test_video_concat_splice.py b/tests/app/invocations/test_video_concat_splice.py index b108ff82d6b..2bb2a80f37e 100644 --- a/tests/app/invocations/test_video_concat_splice.py +++ b/tests/app/invocations/test_video_concat_splice.py @@ -11,6 +11,7 @@ Real ffmpeg (the imageio-ffmpeg bundled binary) encodes and decodes throughout. """ +import importlib import shutil import struct import subprocess @@ -22,7 +23,6 @@ import numpy as np import pytest -from invokeai.app.invocations import video_concat from invokeai.app.invocations.fields import VideoField from invokeai.app.invocations.video_concat import VideoConcatInvocation from invokeai.app.services.session_processor.session_processor_common import CanceledException @@ -385,7 +385,16 @@ def test_first_clip_with_presentation_tags_is_re_encoded(tmp_path, tag): assert len(_decode(out)) == 138 -@pytest.mark.parametrize("stage", ["copy_frames", "splice_soundtrack"]) +@pytest.mark.parametrize( + "stage", + [ + "invokeai.app.invocations.video_concat.copy_frames", + "invokeai.app.invocations.video_concat.splice_soundtrack", + # The audio decodes: an input's soundtrack, and the splice's own validation decodes. + "invokeai.app.invocations.video_concat.extract_audio_pcm", + "invokeai.app.util.aac_splice.extract_audio_pcm", + ], +) def test_cancellation_during_a_splice_saves_nothing(tmp_path, stage): """Cancelling while a splice stage runs stops it — the stage itself raises, not just the pre-save check — and nothing is saved.""" @@ -393,7 +402,8 @@ def test_cancellation_during_a_splice_saves_nothing(tmp_path, stage): b = _encode(tmp_path / "b.mp4", _frames(40, seed=2), tone=660.0) canceled = [False] raised: list[BaseException] = [] - real = getattr(video_concat, stage) + module, name = stage.rsplit(".", 1) + real = getattr(importlib.import_module(module), name) def cancel_then_run(*args, **kwargs): canceled[0] = True @@ -407,7 +417,7 @@ def cancel_then_run(*args, **kwargs): context.videos.get_path.side_effect = lambda name: {"a": a, "b": b}[name] context.util.is_canceled.side_effect = lambda: canceled[0] node = VideoConcatInvocation(videos=[VideoField(video_name="a"), VideoField(video_name="b")], transition="cut") - with patch(f"invokeai.app.invocations.video_concat.{stage}", side_effect=cancel_then_run): + with patch(stage, side_effect=cancel_then_run): with pytest.raises(CanceledException): node.invoke(context) assert len(raised) == 1 and isinstance(raised[0], CanceledException) diff --git a/tests/app/util/test_h264_splice.py b/tests/app/util/test_h264_splice.py index 4ec5191a796..3a4aa1daae6 100644 --- a/tests/app/util/test_h264_splice.py +++ b/tests/app/util/test_h264_splice.py @@ -6,6 +6,7 @@ by a stream copy; either would misalign or misorient a join. Real ffmpeg throughout. """ +import io import os import struct import subprocess @@ -201,13 +202,23 @@ def spawn(*args, **kwargs) -> subprocess.Popen: assert children and children[-1].returncode is not None, "the child was left running" -def _grow_sample_entry(data: bytes, handler: bytes, blob: bytes) -> bytes: - """Append ``blob`` inside the ``handler`` track's sample entry, growing every ancestor box. - The file must keep ``moov`` after ``mdat`` so no chunk offset moves.""" +def _grow_sample_entry( + data: bytes, + handler: bytes, + blob: bytes, + first: bool = False, + inside: bytes | None = None, + in_stbl: bool = False, +) -> bytes: + """Insert ``blob`` inside the ``handler`` track's sample entry — after its other child boxes, + before them if ``first``, or at the end of its ``inside`` child box — or, with ``in_stbl``, + at the start of the track's sample table, growing every ancestor box. The file must keep + ``moov`` after ``mdat`` so no chunk offset moves.""" def boxes(start: int, end: int): while start + 8 <= end: size, kind = struct.unpack(">I4s", data[start : start + 8]) + assert size >= 8, "the fixture must be a well-formed MP4" yield kind, start, size start += size @@ -224,20 +235,33 @@ def boxes(start: int, end: int): hdlr = data.find(b"hdlr", trak, trak + trak_size) if data[hdlr + 12 : hdlr + 16] != handler: continue + grown = bytearray(data) + if in_stbl: + # Ahead of stsd and every other table. + chain.pop() + stbl_at = chain[-1][0] + grown[stbl_at + 8 : stbl_at + 8] = blob + for box_at, size in chain: + struct.pack_into(">I", grown, box_at, size + len(blob)) + return bytes(grown) _, entry, entry_size = next(boxes(start + 8, end)) chain.append((entry, entry_size)) - grown = bytearray(data) - grown[entry + entry_size : entry + entry_size] = blob + # Child boxes follow the entry's box header and fixed fields: 78 bytes for avc1, 28 for mp4a. + children_at = entry + 8 + (78 if handler == b"vide" else 28) + at = children_at if first else entry + entry_size + if inside is not None: + _, child, child_size = next(b for b in boxes(children_at, entry + entry_size) if b[0] == inside) + chain.append((child, child_size)) + at = child + child_size + grown[at:at] = blob for at, size in chain: struct.pack_into(">I", grown, at, size + len(blob)) return bytes(grown) raise AssertionError(f"no {handler!r} track") -@pytest.mark.parametrize("handler", [b"vide", b"soun"]) -def test_sample_entry_stuffed_with_boxes_is_refused_cheaply(tmp_path: Path, handler: bytes) -> None: - """The sample entry is the uploader's: a flood of child boxes must be refused after a - bounded look, not collected (memory) or concatenated (quadratic time).""" +def _audio_clip(tmp_path: Path) -> Path: + """A clip of this app's own encoding, with a soundtrack, and moov after mdat.""" wav = tmp_path / "a.wav" write_stereo_wav(wav, np.zeros((2, 48000), dtype=np.float32), 48000) source = tmp_path / "a.mp4" @@ -245,6 +269,14 @@ def test_sample_entry_stuffed_with_boxes_is_refused_cheaply(tmp_path: Path, hand for _ in range(16): writer.append_data(np.zeros((64, 96, 3), dtype=np.uint8)) writer.close() + return source + + +@pytest.mark.parametrize("handler", [b"vide", b"soun"]) +def test_sample_entry_stuffed_with_boxes_is_refused_cheaply(tmp_path: Path, handler: bytes) -> None: + """The sample entry is the uploader's: a flood of child boxes must be refused after a + bounded look, not collected (memory) or concatenated (quadratic time).""" + source = _audio_clip(tmp_path) hostile = tmp_path / "hostile.mp4" hostile.write_bytes(_grow_sample_entry(source.read_bytes(), handler, struct.pack(">I4s", 8, b"free") * 200_000)) @@ -257,3 +289,76 @@ def test_sample_entry_stuffed_with_boxes_is_refused_cheaply(tmp_path: Path, hand tracemalloc.stop() # 200k boxes are 1.6 MB of file; reading a bounded handful of them allocates almost nothing. assert peak < 1 << 20 + + +@pytest.mark.parametrize( + "handler,inside", [(b"vide", None), (b"vide", b"avcC"), (b"soun", b"esds")], ids=["free", "avcC", "esds"] +) +def test_oversized_sample_entry_is_refused_unread(tmp_path: Path, handler: bytes, inside: bytes | None) -> None: + """One huge child box — a free box, or the decoder config padded out — must be refused by + size, not read into memory.""" + source = _audio_clip(tmp_path) + hostile = tmp_path / "hostile.mp4" + padding = 16 << 20 + blob = bytes(padding) if inside else struct.pack(">I4s", 8 + padding, b"free") + bytes(padding) + hostile.write_bytes(_grow_sample_entry(source.read_bytes(), handler, blob, inside=inside)) + read = read_h264_layout if handler == b"vide" else read_aac_layout + tracemalloc.start() + try: + assert read(hostile) is None + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + assert peak < 1 << 20 + + +def test_box_flood_before_avcc_is_not_scanned(tmp_path: Path) -> None: + """``avcC`` is found among a bounded number of child boxes, so a flood placed before it is + refused after a fixed look rather than walked in full.""" + source = _audio_clip(tmp_path) + hostile = tmp_path / "hostile.mp4" + hostile.write_bytes( + _grow_sample_entry(source.read_bytes(), b"vide", struct.pack(">I4s", 8, b"free") * 2_000_000, first=True) + ) + seeks = [0] + + class CountingFile(io.BufferedReader): + def seek(self, *args): + seeks[0] += 1 + return super().seek(*args) + + def counting_open(path, mode="r", *args, **kwargs): + return CountingFile(io.FileIO(path, "r")) + + # The flood really sits inside the sample entry, ahead of its avcC. + data = hostile.read_bytes() + entry = data.find(b"avc1") + assert entry < data.find(b"\x00\x00\x00\x08free", entry) < data.find(b"avcC", entry) + with patch("invokeai.app.util.h264_splice.open", counting_open, create=True): + assert read_h264_layout(hostile) is None + # Walking the flood would take a seek per box: two million of them. + assert seeks[0] < 10_000 + + +@pytest.mark.parametrize("handler", [b"vide", b"soun"]) +def test_box_flood_in_a_container_is_not_scanned(tmp_path: Path, handler: bytes) -> None: + """Every container walk (moov, trak, mdia, minf, stbl) stops after a fixed number of boxes: + a flood ahead of the sample tables would otherwise be walked once per table looked up.""" + source = _audio_clip(tmp_path) + hostile = tmp_path / "hostile.mp4" + hostile.write_bytes( + _grow_sample_entry(source.read_bytes(), handler, struct.pack(">I4s", 8, b"free") * 200_000, in_stbl=True) + ) + module = "h264_splice" if handler == b"vide" else "aac_splice" + read = read_h264_layout if handler == b"vide" else read_aac_layout + seeks = [0] + + class CountingFile(io.BufferedReader): + def seek(self, *args): + seeks[0] += 1 + return super().seek(*args) + + with patch(f"invokeai.app.util.{module}.open", lambda path, *a, **k: CountingFile(io.FileIO(path)), create=True): + assert read(hostile) is None + # Walking the flood would take a seek per box per lookup: hundreds of thousands at least. + assert seeks[0] < 10_000 From 4b52ab11ab70a7a826226db93374d6d58fce83e0 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Wed, 14 Oct 2026 09:28:28 -0400 Subject: [PATCH 5/7] fix(video): keep whole-range copies to the video's duration The whole-source copy is taken only when no stream plays past the last frame (probed duration), and carries only the first audio track, as a trim does. The RLIMIT_AS allocation test runs only on Linux, which enforces the limit. --- .../invocations/video_frame_extract_range.py | 14 +++- invokeai/app/util/mp4_metadata.py | 16 +++-- ..._video_frame_extract_range_audio_splice.py | 64 +++++++++++++++++++ tests/app/util/test_h264_splice.py | 4 +- 4 files changed, 89 insertions(+), 9 deletions(-) diff --git a/invokeai/app/invocations/video_frame_extract_range.py b/invokeai/app/invocations/video_frame_extract_range.py index c14a8189ee8..545f2dd5572 100644 --- a/invokeai/app/invocations/video_frame_extract_range.py +++ b/invokeai/app/invocations/video_frame_extract_range.py @@ -231,8 +231,15 @@ def invoke(self, context: InvocationContext) -> ExtractVideoRangeOutput: layout = read_h264_layout(video_path) if layout is not None and layout.sample_count != exact_frame_count: layout = None - # The whole clip is the source itself; copying it avoids another lossy generation. - if layout is not None and start == 0 and end == layout.sample_count - 1: + # The whole clip is the source itself; copying it avoids another lossy generation. Unless + # some stream plays past the last frame — audio that outlasts the video or starts late + # (the probed duration is the longest stream's, end trims applied) — which a trim must cut. + if ( + layout is not None + and start == 0 + and end == layout.sample_count - 1 + and duration <= n_frames / output_fps + 0.5 / output_fps + 0.01 + ): return self._save_unchanged(context, video_path, width, height, n_frames, output_fps) context.util.signal_progress(f"Transcoding frames {start}-{end} of {n_frames} @ {output_fps:.2f} fps") @@ -452,7 +459,8 @@ def _save_unchanged( copy_path = Path(work) / "source.mp4" # A stream-copy remux that drops the source's keyed metadata: save() writes this # node's record, and if that rewrite fails the source's must not be read back as it. - write_mp4_tags(video_path, copy_path, {}, context.util.is_canceled) + # Only the first soundtrack: the ordinary path carries no other. + write_mp4_tags(video_path, copy_path, {}, context.util.is_canceled, first_audio_only=True) # A cancellation that arrived during an uninterruptible stage must not be saved. if context.util.is_canceled(): raise CanceledException diff --git a/invokeai/app/util/mp4_metadata.py b/invokeai/app/util/mp4_metadata.py index 1074ff6b75c..b5a8c9f43ec 100644 --- a/invokeai/app/util/mp4_metadata.py +++ b/invokeai/app/util/mp4_metadata.py @@ -215,11 +215,17 @@ def _escape_ffmetadata(value: str) -> str: def write_mp4_tags( - src: Path, dst: Path, tags: Mapping[str, str], is_canceled: Optional[Callable[[], bool]] = None + src: Path, + dst: Path, + tags: Mapping[str, str], + is_canceled: Optional[Callable[[], bool]] = None, + *, + first_audio_only: bool = False, ) -> None: - """Write ``dst`` = ``src``'s first video stream and every audio stream, stream-copied, with - ``tags`` as the file's keyed metadata (replacing any global tags ``src`` had). An empty - mapping therefore strips ``src``'s global tags. + """Write ``dst`` = ``src``'s first video stream and every audio stream (or only the first, + with ``first_audio_only``), stream-copied, with ``tags`` as the file's keyed metadata + (replacing any global tags ``src`` had). An empty mapping therefore strips ``src``'s + global tags. Subtitle, data and further video tracks are not carried: InvokeAI stores only what it serves, and the ingest path already narrows uploads the same way. ``dst`` is written with @@ -251,7 +257,7 @@ def write_mp4_tags( "-map", "0:v:0", "-map", - "0:a?", + "0:a:0?" if first_audio_only else "0:a?", "-map_metadata", "1", "-c", diff --git a/tests/app/invocations/test_video_frame_extract_range_audio_splice.py b/tests/app/invocations/test_video_frame_extract_range_audio_splice.py index d33d649ab8f..9cbad042c1f 100644 --- a/tests/app/invocations/test_video_frame_extract_range_audio_splice.py +++ b/tests/app/invocations/test_video_frame_extract_range_audio_splice.py @@ -12,10 +12,12 @@ """ import struct +import subprocess from pathlib import Path from unittest.mock import MagicMock, patch import imageio.v2 as iio2 +import imageio_ffmpeg import numpy as np import pytest @@ -220,3 +222,65 @@ def test_trim_splices_when_frames_hold_fractional_samples(tmp_path): want = decoded[0][:, round(7 * rate / fps) : round(101 * rate / fps)] interior = slice(LEAD_IN_LIMIT, want.shape[1] - TAIL_LIMIT) assert np.array_equal(got[0][:, interior], want[:, interior]) + + +@pytest.mark.parametrize("n_frames,fps,rate", [(40, 16.0, 32000), (97, 24.0, 32000), (81, 16.0, 16000)]) +def test_whole_clip_copy_keeps_a_soundtrack_that_ends_with_the_video(tmp_path, n_frames, fps, rate): + """The app's own soundtracks decode up to a packet past their end (the decoder does not + apply the end trim); that must not cost them the copy. 97 frames at 24 fps and 32 kHz is + a MiniMax H3 clip, 81 at 16 fps and 16 kHz a 16 kHz Wan-length one.""" + n = round(n_frames / fps * rate) + wav, source = tmp_path / "source.wav", tmp_path / "source.mp4" + write_stereo_wav(wav, np.stack([np.sin(np.arange(n) / 7.0)] * 2).astype(np.float32) * 0.4, rate) + writer = make_mp4_writer(source, fps, audio_path=wav) + for i in range(n_frames): + writer.append_data(np.full((SIZE, SIZE, 3), i, dtype=np.uint8)) + writer.close() + + saved, context = _trim(tmp_path, source, 0, -1) + assert "the whole of the source, copied" in _logged(context) + want, got = extract_audio_pcm(source), extract_audio_pcm(saved) + assert want is not None and got is not None + assert np.array_equal(got[0], want[0]) + + +@pytest.mark.parametrize( + "audio_args,audio_frames,copied", + [ + ([], 80, "Stream-copied frames 0-39"), + (["-itsoffset", "1.0"], 39, "Stream-copied frames 0-39"), + ([], 40, "the whole of the source, copied"), + ], + ids=["longer-audio", "delayed-audio", "second-audio-track"], +) +def test_whole_clip_copy_carries_only_the_soundtrack_a_trim_would(tmp_path, audio_args, audio_frames, copied): + """An uploaded H.264 MP4 is stored as-is, so its audio can outlast its video, start late, + or come as several tracks. Copying the file whole would keep audio past the last frame + (or every track) where any other range is cut to the frames it keeps (and the first track).""" + video = tmp_path / "video.mp4" + writer = make_mp4_writer(video, FPS) + rng = np.random.default_rng(0) + for _ in range(40): + writer.append_data(rng.integers(0, 256, (SIZE, SIZE, 3), dtype=np.uint8)) + writer.close() + audio = tmp_path / "audio.wav" + write_stereo_wav(audio, _signal("tone", audio_frames, 0), RATE) + extra_track = copied.startswith("the whole") + inputs = ["-i", str(video), *audio_args, "-i", str(audio)] + (["-i", str(audio)] if extra_track else []) + maps = ["-map", "0:v", "-map", "1:a"] + (["-map", "2:a"] if extra_track else []) + source = tmp_path / "source.mp4" + subprocess.run( + [imageio_ffmpeg.get_ffmpeg_exe(), "-y", "-loglevel", "error", *inputs, *maps] + + ["-c:v", "copy", "-c:a", "aac", "-aac_pns", "0", str(source)], + check=True, + ) + + saved, context = _trim(tmp_path, source, 0, -1) + # Nothing plays past the last frame (the probe reports to the hundredth of a second). + assert _movie_duration(saved) <= 40 / FPS + 0.5 / FPS + banner = subprocess.run( + [imageio_ffmpeg.get_ffmpeg_exe(), "-hide_banner", "-i", str(saved)], capture_output=True, text=True + ).stderr + assert banner.count("Audio:") == 1 + # The video is still copied, not re-encoded. + assert copied in _logged(context) diff --git a/tests/app/util/test_h264_splice.py b/tests/app/util/test_h264_splice.py index 3a4aa1daae6..1ad8ae25753 100644 --- a/tests/app/util/test_h264_splice.py +++ b/tests/app/util/test_h264_splice.py @@ -118,7 +118,9 @@ def _patch_box(data: bytearray, box: bytes, offset: int, value: bytes) -> None: data[at + 4 + offset : at + 4 + offset + len(value)] = value -@pytest.mark.skipif(sys.platform == "win32", reason="bounds the child's address space with RLIMIT_AS") +@pytest.mark.skipif( + sys.platform != "linux", reason="bounds the child's address space with RLIMIT_AS, which only Linux enforces" +) @pytest.mark.parametrize("uniform_size", [0, 1], ids=["size-table", "uniform-size"]) def test_oversized_sample_count_is_refused_without_allocating(tmp_path: Path, uniform_size: int) -> None: """Uploads that are already H.264 MP4s are stored as-is, so the sample tables are the From 17d06e9d21172b9403d091df61a4e9f83ccd2275 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Wed, 14 Oct 2026 08:52:35 -0500 Subject: [PATCH 6/7] fix(video): guard whole-range copy with movie duration Use the MP4 movie timeline when probe fallback reports only video duration. Cover longer AAC on the OpenCV probe fallback. --- .../invocations/video_frame_extract_range.py | 17 ++++---- invokeai/app/util/mp4_metadata.py | 39 +++++++++++++++++++ ..._video_frame_extract_range_audio_splice.py | 25 ++++++++---- 3 files changed, 66 insertions(+), 15 deletions(-) diff --git a/invokeai/app/invocations/video_frame_extract_range.py b/invokeai/app/invocations/video_frame_extract_range.py index 545f2dd5572..c04c23090e7 100644 --- a/invokeai/app/invocations/video_frame_extract_range.py +++ b/invokeai/app/invocations/video_frame_extract_range.py @@ -65,7 +65,7 @@ probes_as, read_h264_layout, ) -from invokeai.app.util.mp4_metadata import write_mp4_tags +from invokeai.app.util.mp4_metadata import read_mp4_movie_duration, write_mp4_tags from invokeai.app.util.video_audio import extract_audio_pcm, mux_audio_into_video, resample_linear from invokeai.app.util.video_encoding import make_mp4_writer, write_stereo_wav, writer_sample_entry from invokeai.app.util.video_thumbnails import decoder_frame_count, iter_video_frames, probe_video @@ -231,14 +231,15 @@ def invoke(self, context: InvocationContext) -> ExtractVideoRangeOutput: layout = read_h264_layout(video_path) if layout is not None and layout.sample_count != exact_frame_count: layout = None - # The whole clip is the source itself; copying it avoids another lossy generation. Unless - # some stream plays past the last frame — audio that outlasts the video or starts late - # (the probed duration is the longest stream's, end trims applied) — which a trim must cut. + # The whole clip is the source itself; copying it avoids another lossy generation. The + # OpenCV probe fallback knows only the video duration, so also require the MP4 movie header + # to fit the selected frames before bypassing the ordinary path that trims the soundtrack. + whole_range = layout is not None and start == 0 and end == layout.sample_count - 1 + movie_duration = read_mp4_movie_duration(video_path) if whole_range else None if ( - layout is not None - and start == 0 - and end == layout.sample_count - 1 - and duration <= n_frames / output_fps + 0.5 / output_fps + 0.01 + whole_range + and movie_duration is not None + and max(duration, movie_duration) <= n_frames / output_fps + 0.5 / output_fps + 0.01 ): return self._save_unchanged(context, video_path, width, height, n_frames, output_fps) diff --git a/invokeai/app/util/mp4_metadata.py b/invokeai/app/util/mp4_metadata.py index b5a8c9f43ec..d8364850734 100644 --- a/invokeai/app/util/mp4_metadata.py +++ b/invokeai/app/util/mp4_metadata.py @@ -38,6 +38,8 @@ # fragmented file, so a raw upload padded with tiny top-level boxes is given up on rather than # walked. Children of ``moov`` are bounded by ``moov``'s own size. _MAX_TOP_LEVEL_BOXES = 1024 +# Movie headers are near the start of ``moov``; bound this walk for uploaded files. +_MAX_MOVIE_HEADER_CHILDREN = 64 # Same wall-clock bound as the audio mux: a stream copy of a 1 GB upload is disk-bound and # finishes in seconds, so this only stops a wedged child from holding the save forever. REMUX_TIMEOUT_SECONDS = 600 @@ -92,6 +94,43 @@ def find_box( return None +def read_mp4_movie_duration(path: Path) -> Optional[float]: + """Return the MP4 movie-header duration in seconds, or None if it is unavailable or invalid. + + Unlike OpenCV's video-only fallback probe, ``mvhd`` covers the movie timeline across tracks. + The bounded box walks and fixed-size reads keep this safe for uploaded containers. + """ + with open(path, "rb") as fh: + fh.seek(0, 2) + moov = find_box(fh, 0, fh.tell(), b"moov", max_boxes=_MAX_TOP_LEVEL_BOXES) + if moov is None: + return None + mvhd = find_box(fh, *moov, b"mvhd", max_boxes=_MAX_MOVIE_HEADER_CHILDREN) + if mvhd is None or mvhd[0] >= mvhd[1]: + return None + fh.seek(mvhd[0]) + version = fh.read(1) + if not version or version[0] not in (0, 1): + return None + if version[0] == 0: + timescale_offset, duration_offset, duration_size = 12, 16, 4 + else: + timescale_offset, duration_offset, duration_size = 20, 24, 8 + if mvhd[0] + duration_offset + duration_size > mvhd[1]: + return None + fh.seek(mvhd[0] + timescale_offset) + timescale_raw = fh.read(4) + fh.seek(mvhd[0] + duration_offset) + duration_raw = fh.read(duration_size) + if len(timescale_raw) != 4 or len(duration_raw) != duration_size: + return None + timescale = struct.unpack(">I", timescale_raw)[0] + duration = int.from_bytes(duration_raw, "big") + if timescale == 0 or duration == (1 << (duration_size * 8)) - 1: + return None + return duration / timescale + + def _meta_payload_start(fh: BinaryIO, start: int, end: int) -> int: """``meta`` is a FullBox (4 bytes of version/flags) in ISO BMFF but a plain box in QuickTime files. Disambiguate the way ffmpeg does: if the bytes right after the header diff --git a/tests/app/invocations/test_video_frame_extract_range_audio_splice.py b/tests/app/invocations/test_video_frame_extract_range_audio_splice.py index 9cbad042c1f..0cb344d5227 100644 --- a/tests/app/invocations/test_video_frame_extract_range_audio_splice.py +++ b/tests/app/invocations/test_video_frame_extract_range_audio_splice.py @@ -245,17 +245,21 @@ def test_whole_clip_copy_keeps_a_soundtrack_that_ends_with_the_video(tmp_path, n @pytest.mark.parametrize( - "audio_args,audio_frames,copied", + "audio_args,audio_frames,copied,video_only_probe", [ - ([], 80, "Stream-copied frames 0-39"), - (["-itsoffset", "1.0"], 39, "Stream-copied frames 0-39"), - ([], 40, "the whole of the source, copied"), + ([], 80, "Stream-copied frames 0-39", False), + (["-itsoffset", "1.0"], 39, "Stream-copied frames 0-39", False), + ([], 40, "the whole of the source, copied", False), + ([], 80, "Stream-copied frames 0-39", True), ], - ids=["longer-audio", "delayed-audio", "second-audio-track"], + ids=["longer-audio", "delayed-audio", "second-audio-track", "longer-audio-opencv-probe"], ) -def test_whole_clip_copy_carries_only_the_soundtrack_a_trim_would(tmp_path, audio_args, audio_frames, copied): +def test_whole_clip_copy_carries_only_the_soundtrack_a_trim_would( + tmp_path, monkeypatch, audio_args, audio_frames, copied, video_only_probe +): """An uploaded H.264 MP4 is stored as-is, so its audio can outlast its video, start late, - or come as several tracks. Copying the file whole would keep audio past the last frame + or come as several tracks. The OpenCV probe fallback knows only the video duration, so the + MP4 movie duration must also guard a whole-file copy that would keep audio past its last frame (or every track) where any other range is cut to the frames it keeps (and the first track).""" video = tmp_path / "video.mp4" writer = make_mp4_writer(video, FPS) @@ -275,6 +279,13 @@ def test_whole_clip_copy_carries_only_the_soundtrack_a_trim_would(tmp_path, audi check=True, ) + if video_only_probe: + # video_decode_worker._probe() returns frame_count / fps when it falls back to OpenCV. + monkeypatch.setattr( + "invokeai.app.invocations.video_frame_extract_range.probe_video", + lambda _path: (SIZE, SIZE, 40 / FPS, FPS), + ) + saved, context = _trim(tmp_path, source, 0, -1) # Nothing plays past the last frame (the probe reports to the hundredth of a second). assert _movie_duration(saved) <= 40 / FPS + 0.5 / FPS From 76ddfed1c05c97a50dd045f42ec36ab3be834bba Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Wed, 14 Oct 2026 09:56:19 -0500 Subject: [PATCH 7/7] test(video): exercise OpenCV probe fallback Rename the duration-stub regression and cover the real worker fallback. --- ..._video_frame_extract_range_audio_splice.py | 115 +++++++++++++----- 1 file changed, 86 insertions(+), 29 deletions(-) diff --git a/tests/app/invocations/test_video_frame_extract_range_audio_splice.py b/tests/app/invocations/test_video_frame_extract_range_audio_splice.py index 0cb344d5227..7448b7e912f 100644 --- a/tests/app/invocations/test_video_frame_extract_range_audio_splice.py +++ b/tests/app/invocations/test_video_frame_extract_range_audio_splice.py @@ -13,6 +13,7 @@ import struct import subprocess +import sys from pathlib import Path from unittest.mock import MagicMock, patch @@ -24,6 +25,7 @@ from invokeai.app.invocations.fields import VideoField from invokeai.app.invocations.video_concat import VideoConcatInvocation from invokeai.app.invocations.video_frame_extract_range import ExtractVideoRangeInvocation +from invokeai.app.util import video_thumbnails from invokeai.app.util.aac_splice import EDGE_FRAMES, PREROLL_FRAMES, splice_soundtrack, spliceable_aac_layout from invokeai.app.util.mp4_metadata import find_box, write_mp4_tags from invokeai.app.util.video_audio import extract_audio_pcm @@ -61,6 +63,36 @@ def _clip(dirpath: Path, name: str, n_frames: int, kind: str, seed: int) -> Path return path +def _uploaded_clip_with_audio( + tmp_path: Path, + *, + audio_args: list[str], + audio_frames: int, + second_audio_track: bool = False, +) -> Path: + video = tmp_path / "video.mp4" + writer = make_mp4_writer(video, FPS) + rng = np.random.default_rng(0) + for _ in range(40): + writer.append_data(rng.integers(0, 256, (SIZE, SIZE, 3), dtype=np.uint8)) + writer.close() + + audio = tmp_path / "audio.wav" + write_stereo_wav(audio, _signal("tone", audio_frames, 0), RATE) + inputs = ["-i", str(video), *audio_args, "-i", str(audio)] + maps = ["-map", "0:v", "-map", "1:a"] + if second_audio_track: + inputs += ["-i", str(audio)] + maps += ["-map", "2:a"] + source = tmp_path / "source.mp4" + subprocess.run( + [imageio_ffmpeg.get_ffmpeg_exe(), "-y", "-loglevel", "error", *inputs, *maps] + + ["-c:v", "copy", "-c:a", "aac", "-aac_pns", "0", str(source)], + check=True, + ) + return source + + def _save_to(saved: Path): def _save(source_path: Path, **_: object) -> MagicMock: # The stream-copy remux save() runs to embed the record. @@ -245,42 +277,28 @@ def test_whole_clip_copy_keeps_a_soundtrack_that_ends_with_the_video(tmp_path, n @pytest.mark.parametrize( - "audio_args,audio_frames,copied,video_only_probe", + "audio_args,audio_frames,copied,second_audio_track,video_only_probe", [ - ([], 80, "Stream-copied frames 0-39", False), - (["-itsoffset", "1.0"], 39, "Stream-copied frames 0-39", False), - ([], 40, "the whole of the source, copied", False), - ([], 80, "Stream-copied frames 0-39", True), + ([], 80, "Stream-copied frames 0-39", False, False), + (["-itsoffset", "1.0"], 39, "Stream-copied frames 0-39", False, False), + ([], 40, "the whole of the source, copied", True, False), + ([], 80, "Stream-copied frames 0-39", False, True), ], - ids=["longer-audio", "delayed-audio", "second-audio-track", "longer-audio-opencv-probe"], + ids=["longer-audio", "delayed-audio", "second-audio-track", "longer-audio-video-only-probe-result"], ) -def test_whole_clip_copy_carries_only_the_soundtrack_a_trim_would( - tmp_path, monkeypatch, audio_args, audio_frames, copied, video_only_probe +def test_whole_clip_copy_does_not_trust_a_video_only_probe_duration( + tmp_path, monkeypatch, audio_args, audio_frames, copied, second_audio_track, video_only_probe ): - """An uploaded H.264 MP4 is stored as-is, so its audio can outlast its video, start late, - or come as several tracks. The OpenCV probe fallback knows only the video duration, so the - MP4 movie duration must also guard a whole-file copy that would keep audio past its last frame - (or every track) where any other range is cut to the frames it keeps (and the first track).""" - video = tmp_path / "video.mp4" - writer = make_mp4_writer(video, FPS) - rng = np.random.default_rng(0) - for _ in range(40): - writer.append_data(rng.integers(0, 256, (SIZE, SIZE, 3), dtype=np.uint8)) - writer.close() - audio = tmp_path / "audio.wav" - write_stereo_wav(audio, _signal("tone", audio_frames, 0), RATE) - extra_track = copied.startswith("the whole") - inputs = ["-i", str(video), *audio_args, "-i", str(audio)] + (["-i", str(audio)] if extra_track else []) - maps = ["-map", "0:v", "-map", "1:a"] + (["-map", "2:a"] if extra_track else []) - source = tmp_path / "source.mp4" - subprocess.run( - [imageio_ffmpeg.get_ffmpeg_exe(), "-y", "-loglevel", "error", *inputs, *maps] - + ["-c:v", "copy", "-c:a", "aac", "-aac_pns", "0", str(source)], - check=True, + """The invocation must trim out-of-range audio even when its probe reports only video duration.""" + source = _uploaded_clip_with_audio( + tmp_path, + audio_args=audio_args, + audio_frames=audio_frames, + second_audio_track=second_audio_track, ) if video_only_probe: - # video_decode_worker._probe() returns frame_count / fps when it falls back to OpenCV. + # Simulate the duration returned by video_decode_worker._probe()'s OpenCV fallback. monkeypatch.setattr( "invokeai.app.invocations.video_frame_extract_range.probe_video", lambda _path: (SIZE, SIZE, 40 / FPS, FPS), @@ -295,3 +313,42 @@ def test_whole_clip_copy_carries_only_the_soundtrack_a_trim_would( assert banner.count("Audio:") == 1 # The video is still copied, not re-encoded. assert copied in _logged(context) + + +def test_whole_clip_longer_audio_trims_when_opencv_probe_fallback_runs(tmp_path, monkeypatch): + """Force imageio's probe to fail in the worker and exercise its real OpenCV fallback.""" + source = _uploaded_clip_with_audio(tmp_path, audio_args=[], audio_frames=80) + worker_wrapper = tmp_path / "force_opencv_probe.py" + worker_wrapper.write_text( + "import imageio.v3 as iio\n" + "import runpy\n" + "import sys\n" + "\n" + "def fail_immeta(*_args, **_kwargs):\n" + " raise RuntimeError('force the OpenCV fallback')\n" + "\n" + "iio.immeta = fail_immeta\n" + "worker_path = sys.argv[1]\n" + "sys.argv = [worker_path, *sys.argv[2:]]\n" + "runpy.run_path(worker_path, run_name='__main__')\n" + ) + original_worker_command = video_thumbnails._worker_command + worker_path = original_worker_command("probe", str(source))[1] + + def worker_command(*args: str) -> list[str]: + if args[0] == "probe": + return [sys.executable, str(worker_wrapper), worker_path, *args] + return original_worker_command(*args) + + monkeypatch.setattr(video_thumbnails, "_worker_command", worker_command) + + width, height, probed_duration, fps = video_thumbnails.probe_video(source) + assert (width, height) == (SIZE, SIZE) + assert probed_duration == pytest.approx(40 / FPS, abs=0.1) + assert fps == pytest.approx(FPS) + assert _movie_duration(source) > 4.9 + + saved, context = _trim(tmp_path, source, 0, -1) + + assert _movie_duration(saved) <= 40 / FPS + 0.5 / FPS + assert "Stream-copied frames 0-39" in _logged(context)