fix(video): Concatenate Videos preserves audio (cut and transitions)

invokeai/app/invocations/video_concat.py

@@ -10,6 +10,14 @@
 at a time, buffering only the transition windows, so peak memory stays
 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 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
+frames, so even upload-cap-length inputs stay within a few hundred MB — but it
+is O(total duration), not O(transition_frames). Audio problems fail the node
+loudly rather than emit another silent video: a silent output masquerading as
+success is the exact bug this path exists to fix.
 """
 
 import math
@@ -30,7 +38,8 @@
 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_encoding import make_mp4_writer
+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_thumbnails import iter_video_frames, probe_video
 
 TransitionMode = Literal["cut", "crossfade", "fade_through_black"]
@@ -69,12 +78,19 @@ def _fade_through_black(a_tail: list[np.ndarray], b_head: list[np.ndarray]) -> I
     title="Concatenate Videos",
     tags=["video", "concat", "transition"],
     category="video",
-    version="1.0.0",
+    version="1.1.0",
     classification=Classification.Prototype,
 )
 class VideoConcatInvocation(BaseInvocation, WithMetadata, WithBoard):
     """Join two or more videos into a single MP4.
 
+    Audio: if any input carries an audio track, the output gets an AAC track assembled on
+    the emitted timeline — silent inputs contribute silence, `cut` splices the tracks,
+    `crossfade` blends them equal-power over the transition window, and
+    `fade_through_black` fades the outgoing track to silence and the incoming one up from
+    it, mirroring the video. An fps override retimes audio with the video (a deliberate
+    speed/pitch change). If no input has audio, the output has no audio stream.
+
     Transitions:
 
     * ``cut`` — hard splice, no blending. Fastest; total length is the sum of inputs.
@@ -143,14 +159,19 @@ def invoke(self, context: InvocationContext) -> VideoOutput:
         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:
                 clip_iters = [iter_video_frames(p, is_canceled=context.util.is_canceled) for p in paths]
-                for frame in self._iter_joined_frames(clip_iters, is_canceled=context.util.is_canceled):
+                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:
@@ -159,13 +180,39 @@ def invoke(self, context: InvocationContext) -> VideoOutput:
             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)
+                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
+
             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}"
+                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=tmp_path,
+                source_path=source_path,
                 width=width,
                 height=height,
                 duration=duration,
@@ -174,10 +221,13 @@ def invoke(self, context: InvocationContext) -> VideoOutput:
             context.logger.info(f"Saved concatenated video: {video_dto.video_name}")
             return VideoOutput.build(video_dto)
         finally:
-            try:
-                tmp_path.unlink(missing_ok=True)
-            except Exception:
-                pass
+            for cleanup in (tmp_path, wav_path, muxed_path):
+                if cleanup is None:
+                    continue
+                try:
+                    cleanup.unlink(missing_ok=True)
+                except Exception:
+                    pass
 
     def _estimate_transition_memory(self, width: int, height: int) -> int:
         if self.transition == "cut" or self.transition_frames == 0:
@@ -212,10 +262,124 @@ def _validate_transition_memory(self, width: int, height: int) -> None:
                 "Lower transition_frames or use lower-resolution clips."
             )
 
+    def _build_audio_track(
+        self,
+        context: InvocationContext,
+        paths: list[Path],
+        frame_counts: list[int],
+        native_rates: list[Optional[float]],
+        native_durations: list[Optional[float]],
+        output_fps: float,
+        num_output_frames: int,
+    ) -> Optional[tuple[np.ndarray, int]]:
+        """Assemble the output soundtrack on the emitted timeline.
+
+        Returns ``(pcm, sample_rate)`` — ``pcm`` shaped ``(2, n)``, float in [-1, 1] — or
+        ``None`` when no input carries audio. The frame bookkeeping (which frames of each
+        clip are emitted directly vs consumed into a boundary) mirrors
+        ``_iter_joined_frames`` exactly; ``frame_counts`` must be that pass's per-clip
+        decoded frame counts so the audio is cut against the frames actually emitted.
+
+        Every clip's audio is mapped onto the emitted timeline with one linear resample:
+        the clip's video spans ``n_i / native_fps`` seconds natively and ``n_i /
+        output_fps`` seconds in the output, so this single step covers both sample-rate
+        unification (to the first audible clip's rate) and any fps-override retime.
+        """
+        context.util.signal_progress("Extracting audio")
+        extracted: list[Optional[tuple[np.ndarray, int]]] = []
+        for path in paths:
+            if context.util.is_canceled():
+                raise CanceledException
+            extracted.append(extract_audio_pcm(path))
+        if all(item is None for item in extracted):
+            return None
+        if context.util.is_canceled():
+            raise CanceledException
+        rate = next(item[1] for item in extracted if item is not None)
+
+        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
+
+        clips: list[np.ndarray] = []
+        for i, (item, n_frames, native_fps) in enumerate(zip(extracted, frame_counts, native_rates, strict=True)):
+            want = s(n_frames)
+            if item is None:
+                clips.append(np.zeros((2, want), dtype=np.float32))
+                continue
+            pcm, src_rate = item
+            # The clip's audio must be sliced to the span its video occupies natively
+            # (trimming AAC end-padding) before the resample maps it onto the emitted
+            # timeline. Preference order for that native span: the probed fps; the
+            # probed duration (fps can be None for VFR-flagged containers, and under an
+            # fps override guessing output_fps would skip the retime — the audio would
+            # play at the wrong speed and truncate or pad); the extracted length itself,
+            # which approximates the native span to within the codec padding.
+            if native_fps is not None and native_fps > 0:
+                src_want = round(n_frames / native_fps * src_rate)
+            elif native_durations[i] is not None and native_durations[i] > 0:
+                src_want = round(native_durations[i] * src_rate)
+            else:
+                src_want = pcm.shape[1]
+            if pcm.shape[1] < src_want:
+                pcm = np.pad(pcm, ((0, 0), (0, src_want - pcm.shape[1])))
+            else:
+                pcm = pcm[:, :src_want]
+            clips.append(resample_linear(pcm, want))
+            extracted[i] = None  # drop the pre-resample copy; peak memory is O(total audio)
+
+        # A boundary exists between consecutive clips whenever the transition consumes
+        # any frames — even when only ONE side contributes (fade_through_black with
+        # transition_frames=1 has an empty tail window but still fades the head in).
+        has_boundary = (head_need + tail_need) > 0
+
+        segments: list[np.ndarray] = []
+        prev_tail: Optional[np.ndarray] = None
+        for i, pcm in enumerate(clips):
+            n_i = frame_counts[i]
+            head_want = 0 if i == 0 else head_need
+            tail_keep = 0 if i == len(clips) - 1 else tail_need
+            head_end = s(head_want)
+            tail_start = s(n_i - tail_keep) if tail_keep else pcm.shape[1]
+            if prev_tail is not None:
+                head = pcm[:, :head_end]
+                if self.transition == "crossfade":
+                    # Equal-power blend: constant perceived loudness through the overlap.
+                    n = min(prev_tail.shape[1], head.shape[1])
+                    theta = np.linspace(0.0, np.pi / 2.0, n, dtype=np.float32)
+                    segments.append(prev_tail[:, :n] * np.cos(theta) + head[:, :n] * np.sin(theta))
+                else:  # fade_through_black: out to silence, then in from silence.
+                    segments.append(prev_tail * np.linspace(1.0, 0.0, prev_tail.shape[1], dtype=np.float32))
+                    segments.append(head * np.linspace(0.0, 1.0, head.shape[1], dtype=np.float32))
+            segments.append(pcm[:, head_end:tail_start])
+            # An empty tail window (shape (2, 0)) still marks the boundary as pending so
+            # the next clip's head is faded in rather than silently dropped.
+            prev_tail = pcm[:, tail_start:] if (has_boundary and i < len(clips) - 1) else None
+
+        pcm_out = np.concatenate(segments, axis=1)
+        # Per-clip rounding can drift the total by a sample per boundary; pin the track
+        # to the emitted video duration exactly.
+        target = s(num_output_frames)
+        if pcm_out.shape[1] < target:
+            pcm_out = np.pad(pcm_out, ((0, 0), (0, target - pcm_out.shape[1])))
+        else:
+            pcm_out = pcm_out[:, :target]
+        return pcm_out.astype(np.float32, copy=False), rate
+
     def _iter_joined_frames(
         self,
         clips: list[Iterable[np.ndarray]],
         is_canceled: Optional[Callable[[], bool]] = None,
+        frame_counts: Optional[list[int]] = None,
     ) -> Iterator[np.ndarray]:
         """Yields the joined output frames, pulling lazily from each clip's frame iterator.
 
@@ -275,6 +439,8 @@ def _iter_joined_frames(
                 tail_buf.append(frame)
                 if len(tail_buf) > tail_keep:
                     yield tail_buf.popleft()
+            if frame_counts is not None:
+                frame_counts.append(n_frames)
             if n_frames == 0:
                 raise ValueError(f"Input video {i} ({self.videos[i].video_name}) decoded to zero frames.")
             if n_frames < head_want + tail_keep: