feat(video): support muxing an audio track in make_mp4_writer

invokeai/app/util/video_encoding.py

@@ -9,13 +9,60 @@
 checks downstream (e.g. concatenating a trimmed clip with its source).
 
 yuv420p requires even dimensions; callers validate that before 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
+  *max* of the stream durations. Callers must trim the audio to the video
+  duration (num_frames / fps seconds) before handing it over, or the player
+  shows trailing frozen video.
+- ``audio_codec`` must name a real encoder. The default AAC-LC plays in every
+  browser; never pass ``"copy"`` for a PCM WAV — PCM-in-MP4 does not play in
+  browsers.
 """
 
+import wave
 from pathlib import Path
 
 import imageio.v2 as iio2
+import numpy as np
+
+
+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.
+
+    If ``audio_path`` is given, that file's audio is encoded with ``audio_codec``
+    and muxed into the output (see module docstring for the trim requirement).
+    The audio file must exist until at least the first ``append_data`` call has
+    completed: ffmpeg is spawned lazily on the first append, and a file deleted
+    before then fails as an unexplained BrokenPipeError later in the encode (or,
+    for tiny clips, silently produces no output at all). The existence check
+    here converts the common failure into an immediate, named error.
+    """
+    kwargs = {}
+    if audio_path is not None:
+        if not Path(audio_path).exists():
+            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)
+
 
+def write_stereo_wav(path: Path | str, samples: np.ndarray, sample_rate: int) -> None:
+    """Write float PCM as a 16-bit stereo WAV suitable for ``make_mp4_writer``'s ``audio_path``.
 
-def make_mp4_writer(path: Path | str, fps: float):
-    """Returns an imageio FFMPEG writer that preserves frame dimensions exactly."""
-    return iio2.get_writer(str(path), format="FFMPEG", mode="I", fps=fps, codec="libx264", macro_block_size=1)
+    ``samples`` must be shaped ``(2, n_samples)`` (channels first), values in
+    [-1, 1]; anything outside is clipped rather than wrapped, and NaN becomes
+    silence (0) rather than undefined int16 garbage. The conversion is done in
+    float64 regardless of input dtype: in float16, ``1.0 * 32767.0`` rounds up
+    to 32770 and would wrap full-scale peaks to -32768.
+    """
+    if samples.ndim != 2 or samples.shape[0] != 2:
+        raise ValueError(f"Expected samples shaped (2, n_samples), got {samples.shape}")
+    widened = np.nan_to_num(samples.astype(np.float64), nan=0.0)
+    int16 = (np.clip(widened, -1.0, 1.0) * 32767.0).astype(np.int16)
+    with wave.open(str(path), "wb") as wav:
+        wav.setnchannels(2)
+        wav.setsampwidth(2)
+        wav.setframerate(sample_rate)
+        wav.writeframes(int16.T.reshape(-1).tobytes())

tests/app/util/test_video_encoding.py

@@ -15,7 +15,7 @@
 import pytest
 
 from invokeai.app.invocations.video_frame_extract_range import _validate_even_dimensions
-from invokeai.app.util.video_encoding import make_mp4_writer
+from invokeai.app.util.video_encoding import make_mp4_writer, write_stereo_wav
 
 
 def test_non_multiple_of_16_dimensions_are_preserved(tmp_path: Path) -> None:
@@ -38,6 +38,124 @@ def test_non_multiple_of_16_dimensions_are_preserved(tmp_path: Path) -> None:
     assert first.shape[:2] == (height, width)
 
 
+def _write_sine_wav(path: Path, duration_s: float, sample_rate: int = 32_000) -> None:
+    t = np.arange(int(duration_s * sample_rate)) / sample_rate
+    samples = np.stack([np.sin(2 * np.pi * 440 * t), np.sin(2 * np.pi * 660 * t)])
+    write_stereo_wav(path, samples, sample_rate)
+
+
+def test_audio_path_muxes_an_aac_stream(tmp_path: Path) -> None:
+    fps, num_frames = 8.0, 8
+    wav_path = tmp_path / "audio.wav"
+    _write_sine_wav(wav_path, duration_s=num_frames / fps)
+
+    path = tmp_path / "out.mp4"
+    writer = make_mp4_writer(path, fps=fps, audio_path=wav_path)
+    try:
+        for i in range(num_frames):
+            writer.append_data(np.full((84, 120, 3), i * 10, dtype=np.uint8))
+    finally:
+        writer.close()
+
+    reader = iio2.get_reader(str(path))
+    try:
+        meta = reader.get_meta_data()
+    finally:
+        reader.close()
+    assert meta["audio_codec"] == "aac"
+    # The audio was pre-trimmed to the video duration, so the container must not
+    # be stretched past it (ffmpeg gets no -shortest; see module docstring).
+    assert meta["duration"] == pytest.approx(num_frames / fps, abs=0.1)
+    assert meta["size"] == (120, 84)
+
+
+def test_audio_codec_is_forwarded_not_defaulted(tmp_path: Path) -> None:
+    # ffmpeg's default MP4 audio encoder is aac, so the aac test above cannot
+    # distinguish "audio_codec forwarded" from "audio_codec dropped". A
+    # non-default codec can.
+    wav_path = tmp_path / "audio.wav"
+    _write_sine_wav(wav_path, duration_s=0.5)
+
+    path = tmp_path / "out.mp4"
+    writer = make_mp4_writer(path, fps=8.0, audio_path=wav_path, audio_codec="libmp3lame")
+    try:
+        for _ in range(4):
+            writer.append_data(np.zeros((84, 120, 3), dtype=np.uint8))
+    finally:
+        writer.close()
+
+    reader = iio2.get_reader(str(path))
+    try:
+        meta = reader.get_meta_data()
+    finally:
+        reader.close()
+    assert meta["audio_codec"] == "mp3"
+
+
+def test_missing_audio_file_fails_fast(tmp_path: Path) -> None:
+    with pytest.raises(FileNotFoundError, match="does-not-exist.wav"):
+        make_mp4_writer(tmp_path / "out.mp4", fps=8.0, audio_path=tmp_path / "does-not-exist.wav")
+
+
+def test_no_audio_path_produces_no_audio_stream(tmp_path: Path) -> None:
+    path = tmp_path / "out.mp4"
+    writer = make_mp4_writer(path, fps=8.0)
+    try:
+        for _ in range(4):
+            writer.append_data(np.zeros((84, 120, 3), dtype=np.uint8))
+    finally:
+        writer.close()
+
+    reader = iio2.get_reader(str(path))
+    try:
+        meta = reader.get_meta_data()
+    finally:
+        reader.close()
+    assert meta.get("audio_codec") in (None, "")
+
+
+def test_write_stereo_wav_format_and_clipping(tmp_path: Path) -> None:
+    import wave
+
+    sample_rate = 32_000
+    samples = np.stack([np.full(100, 2.0), np.full(100, -2.0)])  # out of range -> clipped
+    path = tmp_path / "a.wav"
+    write_stereo_wav(path, samples, sample_rate)
+
+    with wave.open(str(path), "rb") as wav:
+        assert wav.getnchannels() == 2
+        assert wav.getsampwidth() == 2
+        assert wav.getframerate() == sample_rate
+        assert wav.getnframes() == 100
+        frames = np.frombuffer(wav.readframes(100), dtype=np.int16).reshape(-1, 2)
+    assert frames[:, 0].max() == 32767
+    assert frames[:, 1].min() == -32767
+
+    with pytest.raises(ValueError, match=r"\(2, n_samples\)"):
+        write_stereo_wav(tmp_path / "b.wav", np.zeros((100,)), sample_rate)
+
+
+def test_write_stereo_wav_float16_and_nan_are_safe(tmp_path: Path) -> None:
+    import wave
+
+    # float16: 1.0 * 32767.0 rounds to 32770 in fp16 and would wrap to -32768
+    # if the conversion ran in the input dtype.
+    fp16 = np.ones((2, 8), dtype=np.float16)
+    path = tmp_path / "fp16.wav"
+    write_stereo_wav(path, fp16, 32_000)
+    with wave.open(str(path), "rb") as wav:
+        frames = np.frombuffer(wav.readframes(8), dtype=np.int16)
+    assert frames.min() == 32767
+
+    # NaN becomes silence, not undefined int16 garbage.
+    nan = np.full((2, 8), np.nan)
+    path = tmp_path / "nan.wav"
+    write_stereo_wav(path, nan, 32_000)
+    with wave.open(str(path), "rb") as wav:
+        frames = np.frombuffer(wav.readframes(8), dtype=np.int16)
+    assert np.all(frames == 0)
+
+
 def test_validate_even_dimensions_accepts_even_and_rejects_odd() -> None:
     _validate_even_dimensions(1920, 1080, "ok.mp4")
     with pytest.raises(ValueError, match="even dimensions"):