diff --git a/invokeai/app/util/video_encoding.py b/invokeai/app/util/video_encoding.py index f13aabe2ea8..98bb9a7b284 100644 --- a/invokeai/app/util/video_encoding.py +++ b/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()) diff --git a/tests/app/util/test_video_encoding.py b/tests/app/util/test_video_encoding.py index 14ee177b39f..83b20e1bbc6 100644 --- a/tests/app/util/test_video_encoding.py +++ b/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"):