fix(video): stop re-encoding earlier segments when extending videos

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.

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,20 +284,19 @@ 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",
                     str(ffmeta_path),
                     "-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

invokeai/app/util/video_audio.py

@@ -18,23 +18,77 @@
 
 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
 
     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,27 +199,28 @@ 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:
         stderr = proc.stderr.decode("utf-8", errors="replace")
         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,
     )