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..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. 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 1f9537f4c4a..0d7a48957aa 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, + EDGE_FRAMES, + PREROLL_FRAMES, + splice_soundtrack, + spliceable_aac_layout, +) +from invokeai.app.util.h264_splice import ( + MIN_ENCODED_FRAMES, + H264Layout, + SpliceError, + concat_stream_copy, + copy_frames, + probes_as, + read_h264_layout, +) +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 iter_video_frames, probe_video TransitionMode = Literal["cut", "crossfade", "fade_through_black"] @@ -195,84 +226,275 @@ 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 + 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], output_fps, work_dir, tmp_path ): - 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) + 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 + mux_audio_into_video(tmp_path, wav_path, muxed_path, context.util.is_canceled) + 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") + ) + # 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, + 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.sample_entry != writer_sample_entry(*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, + output_fps: float, + 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.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, context.util.is_canceled) + head = read_h264_layout(head_path) + # 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, 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: + 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, too little precedes the transition, or the splice does not verify. + """ + if first_fps is None or first_fps != output_fps: + return False + 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. + 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], + ) + # 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 def _estimate_transition_memory(self, width: int, height: int) -> int: if self.transition == "cut" or self.transition_frames == 0: @@ -335,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(): @@ -346,14 +568,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 +659,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..c04c23090e7 100644 --- a/invokeai/app/invocations/video_frame_extract_range.py +++ b/invokeai/app/invocations/video_frame_extract_range.py @@ -6,6 +6,15 @@ 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``), 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 the source's). Like ``video_concat``, the audio path buffers the source @@ -14,7 +23,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 +50,24 @@ 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, + SpliceError, + concat_stream_copy, + copy_frames, + probes_as, + read_h264_layout, +) +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 +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 @@ -161,7 +188,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 +223,26 @@ 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. 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 ( + 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) + 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 +251,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: @@ -242,14 +297,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, context.util.is_canceled) source_path = muxed_path out_duration = num_frames / output_fps @@ -257,6 +315,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, @@ -286,6 +347,207 @@ 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.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] + 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.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, context.util.is_canceled) + middle = read_h264_layout(middle_path) + 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, 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: + 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. + # 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 + 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 _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, @@ -311,7 +573,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(): @@ -331,7 +593,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 new file mode 100644 index 00000000000..c46216293a1 --- /dev/null +++ b/invokeai/app/util/aac_splice.py @@ -0,0 +1,503 @@ +"""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 (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 +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, Callable, 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, 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 = 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 +# 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 +# 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. 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 + 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, max_boxes=_MAX_CHILD_BOXES): + if box_type != b"trak": + continue + 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) + 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", 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", max_boxes=_MAX_CHILD_BOXES) + 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] + # 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 + # 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]) + 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", max_boxes=_MAX_CHILD_BOXES) + 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", 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]: + 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 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 + 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_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) + 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]) + 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, max_boxes=_MAX_CHILD_BOXES): + if box_type != b"trak": + continue + 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) + 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", 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]) + wide_edit = fh.read(1)[0] == 1 + fh.seek(elst[0] + 8) + 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 + 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 _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, is_canceled=is_canceled) + 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, + source_skip: int, + first: int, + end: int, + copied_at: int, + timeline: np.ndarray, + sample_rate: int, + out_path: Path, + is_canceled: Optional[Callable[[], bool]] = None, +) -> 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. + """ + 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( + # 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, + ) + 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, is_canceled=is_canceled) + 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/h264_splice.py b/invokeai/app/util/h264_splice.py new file mode 100644 index 00000000000..af35c267a2e --- /dev/null +++ b/invokeai/app/util/h264_splice.py @@ -0,0 +1,544 @@ +"""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* 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. + +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 math +import shutil +import struct +import subprocess +import tempfile +from pathlib import Path +from typing import BinaryIO, Callable, NamedTuple, Optional + +import numpy as np + +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 +# ``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...) 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 +# reorder delay; a shorter part joined to full-delay parts distorts the joined timestamps. +MIN_ENCODED_FRAMES = 3 +# 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) +# 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.""" + 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.""" + + +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, max_boxes=_MAX_CHILD_BOXES): + if box_type != b"trak": + continue + 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", 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", 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", max_boxes=_MAX_CHILD_BOXES)): + return None + 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", max_boxes=_MAX_CHILD_BOXES), 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", max_boxes=_MAX_CHILD_BOXES) + 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", max_boxes=_MAX_CHILD_BOXES) + 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_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: + if uniform_size * sample_count > file_size: + return None + 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", 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)): + return None + sample_count = len(sizes) + fh.seek(stsc[0] + 4) + (run_count,) = struct.unpack(">I", fh.read(4)) + 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 > 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 or per_chunk > sample_count: + return None + for chunk in range(first_chunk - 1, last_chunk): + 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 + 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 + 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: + 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", 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 + + # 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] + # 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]) + decoder_config = fh.read(avcc[1] - avcc[0]) + 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)] + 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", 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. + if sample_count > _MAX_SYNC_ENTRIES: + return None + 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, sizes, sync_samples, file_size) + if locations is None: + return None + 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: + 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, + sample_entry=sample_entry, + frame_duration=frame_duration, + ) + + +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`` stops the child promptly and raises CanceledException (see + ``run_ffmpeg_process``), so a canceled node never waits on a long copy. + """ + 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( + 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 + 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, + is_canceled, + ) + 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, + is_canceled, + ) + segment = Path(work) / "segment1.mp4" + if not segment.exists(): + raise SpliceError(f"ffmpeg wrote no segment {description}") + 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, 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 + 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", + is_canceled, + ) + 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..d8364850734 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 @@ -40,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 @@ -52,13 +52,7 @@ 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]]: +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,15 +85,52 @@ 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 +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 @@ -141,7 +172,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 +200,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 +231,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 @@ -222,10 +253,18 @@ def _escape_ffmetadata(value: str) -> str: ) -def write_mp4_tags(src: Path, dst: Path, tags: Mapping[str, str]) -> 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. +def write_mp4_tags( + 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 (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 @@ -245,12 +284,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", @@ -258,7 +296,7 @@ def write_mp4_tags(src: Path, dst: Path, tags: Mapping[str, str]) -> None: "-map", "0:v:0", "-map", - "0:a?", + "0:a:0?" if first_audio_only else "0:a?", "-map_metadata", "1", "-c", @@ -269,8 +307,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 f7cc6eeb1a9..2a28b259d10 100644 --- a/invokeai/app/util/video_audio.py +++ b/invokeai/app/util/video_audio.py @@ -18,11 +18,26 @@ 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 +# 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: import imageio_ffmpeg @@ -30,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 @@ -45,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: @@ -56,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 @@ -157,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: @@ -171,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", @@ -188,10 +231,9 @@ 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, timeout=600, + is_canceled=is_canceled, ) diff --git a/invokeai/app/util/video_encoding.py b/invokeai/app/util/video_encoding.py index 60c560fcf39..906fcf590f1 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,18 @@ browsers. """ +import tempfile import wave from pathlib import Path 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 + 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 +59,55 @@ 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) + # 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), + format="FFMPEG", + mode="I", + fps=fps, + codec="libx264", + macro_block_size=1, + quality=None, + output_params=["-crf", str(CRF), "-g", str(keyframe_interval), *audio_options], + **kwargs, + ) + + +_writer_sample_entrys: dict[tuple[int, int, float], bytes] = {} + + +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 + — 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_sample_entrys: + config = _probe_sample_entry(width, height, fps) + if config is None: + return None + _writer_sample_entrys[key] = config + return _writer_sample_entrys[key] + + +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: + 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.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/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/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..2bb2a80f37e --- /dev/null +++ b/tests/app/invocations/test_video_concat_splice.py @@ -0,0 +1,443 @@ +"""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 importlib +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.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 +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 + + +@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", + [ + "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.""" + 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] = [] + module, name = stage.rsplit(".", 1) + real = getattr(importlib.import_module(module), name) + + 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(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 c773312c673..1cd2dec8dfd 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,236 @@ 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 + + +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/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..7448b7e912f --- /dev/null +++ b/tests/app/invocations/test_video_frame_extract_range_audio_splice.py @@ -0,0 +1,354 @@ +"""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 +import subprocess +import sys +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.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 +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 _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. + 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]) + + +@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,second_audio_track,video_only_probe", + [ + ([], 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-video-only-probe-result"], +) +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 +): + """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: + # 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), + ) + + 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) + + +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) diff --git a/tests/app/util/test_h264_splice.py b/tests/app/util/test_h264_splice.py new file mode 100644 index 00000000000..1ad8ae25753 --- /dev/null +++ b/tests/app/util/test_h264_splice.py @@ -0,0 +1,366 @@ +"""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 io +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.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 + + +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.sample_entry == writer_sample_entry(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,) + + +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 != "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 + 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, + 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 + + _, 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 + 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)) + # 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") + + +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" + 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() + 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)) + + 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 + + +@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