diff --git a/nodes/multi_video_loader.py b/nodes/multi_video_loader.py index 6fb3f3d..bb12696 100644 --- a/nodes/multi_video_loader.py +++ b/nodes/multi_video_loader.py @@ -5,10 +5,10 @@ # count; Advanced adds filename_N. video_N is lazy (VideoFromFile — decodes nothing) unless # force_rate retimes the clip, which requires decoding its frames. audio_N is decoded # separately with PyAV so audio is available even on the cheap (force_rate == 0) path; a clip -# with no audio track yields None there, same as an unplugged OPTIONAL input. The audio -# decoder mirrors nodes/multi_audio_loader.py's _load_audio (ComfyUI, GPL-3.0 — same license -# as this pack); it is duplicated here rather than imported for the same comfy-free reason -# _input_path is duplicated across the loaders. +# with no (decodable) audio track yields None there, same as an unplugged OPTIONAL input. The +# audio decoder mirrors nodes/multi_audio_loader.py's _load_audio (ComfyUI, GPL-3.0 — same +# license as this pack); it is duplicated here rather than imported for the same comfy-free +# reason _input_path is duplicated across the loaders. import os from fractions import Fraction @@ -42,12 +42,24 @@ def _f32_pcm(wav): def _decode_audio_track(path): """Decode only the audio stream -> {"waveform": Tensor[1,C,S], "sample_rate": native_rate}, - or None if the file has no audio track. Raises only if the file itself can't be opened. + or None if the file has no (decodable) audio track. Raises only if the file itself can't + be opened. + + Stream selection mirrors ComfyUI's own last_decodable_audio_stream + (comfy_api.latest._input_impl.video_types): pick the LAST audio stream FFmpeg can actually + decode, not just streams.audio[0]. A stream FFmpeg has no decoder for (e.g. an iPhone APAC + spatial-audio track) has codec_context = None; forcing index 0 on such a file used to make + this node fail even though the force_rate>0 path — which goes through ComfyUI's own + get_components() and applies this same selection — tolerated it fine. """ with av.open(path) as af: - if not af.streams.audio: + stream = next( + (s for s in reversed(af.streams.audio) + if s.codec_context is not None and s.codec_context.sample_rate), + None, + ) + if stream is None: return None - stream = af.streams.audio[0] sr = stream.codec_context.sample_rate frames = [] for frame in af.decode(streams=stream.index): @@ -67,8 +79,10 @@ def _load_video(path, force_rate, pos): ``pos`` is the file's 1-based position in files_json, included in error messages so a failure can be traced back to a specific row in the UI. ``force_rate == 0`` is the cheap path: video_N is a lazy VideoFromFile (decodes nothing) and audio_N is decoded on its own - via PyAV. ``force_rate > 0`` decodes the video's components once, resamples the frames, - and reuses the already-decoded audio instead of decoding it a second time. + via PyAV. ``force_rate > 0`` normally decodes the video's components once, resamples the + frames, and reuses the already-decoded audio instead of decoding it a second time — unless + the clip is already at the requested rate, in which case it takes the same cheap path as + force_rate == 0 (see the get_frame_rate() check below). """ try: video = InputImpl.VideoFromFile(path) @@ -76,6 +90,15 @@ def _load_video(path, force_rate, pos): audio = _decode_audio_track(path) return video, audio + if abs(float(video.get_frame_rate()) - force_rate) < 1e-6: + # Cheap metadata-only check (get_frame_rate() reads container metadata, no frame + # decoding): the clip is already at the requested rate, so take the same free path + # as force_rate == 0 instead of paying for get_components() just to discard the + # result. audio_N still has to come from _decode_audio_track — nothing above filled + # it in on this branch. + audio = _decode_audio_track(path) + return video, audio + components = video.get_components() src_count = components.images.shape[0] indices = resample_frame_indices(src_count, float(components.frame_rate), force_rate) @@ -87,7 +110,10 @@ def _load_video(path, force_rate, pos): Types.VideoComponents( images=resampled_images, audio=components.audio, - frame_rate=Fraction(force_rate), + # limit_denominator(1001) turns NTSC-style rates (e.g. 29.97) into the clean + # 2997/100 rational instead of Fraction(force_rate)'s exact binary-float + # fraction (8437190785180631/281474976710656). + frame_rate=Fraction(force_rate).limit_denominator(1001), ), bit_depth=video.get_bit_depth(), ) diff --git a/tests/test_multi_video_loader.py b/tests/test_multi_video_loader.py index a102e15..2ddb1e6 100644 --- a/tests/test_multi_video_loader.py +++ b/tests/test_multi_video_loader.py @@ -14,17 +14,31 @@ def input_dir(tmp_path, monkeypatch): return tmp_path -def _write_video(path, fps=24, seconds=1, color=(200, 100, 50)): - """Write a tiny synthetic solid-color video (no audio track) with PyAV.""" +def _write_video(path, fps=24, seconds=1, color=(200, 100, 50), audio_hz=None, audio_rate=44100): + """Write a tiny synthetic solid-color video with PyAV. + + ``audio_hz`` is optional: when given, a mono ``audio_rate``-Hz sine tone of the same + duration is muxed in as an AAC track (the whole waveform is handed to + ``AudioStream.encode()`` in one call — it internally chunks to the codec's frame size, same + pattern ComfyUI's own VideoFromComponents.save_to uses). Left ``None``, the file has no + audio track at all, as before. + """ import av import numpy as np n_frames = int(round(fps * seconds)) container = av.open(str(path), mode="w") + # Both streams must be registered before any packet is muxed — the container writes its + # header on the first mux() call, so an audio stream added only after the video track is + # fully flushed produces a broken/undersized header ("Cannot rebase to zero time" on mux). stream = container.add_stream("libx264", rate=fps) stream.width = 32 stream.height = 32 stream.pix_fmt = "yuv420p" + audio_stream = None + if audio_hz is not None: + audio_stream = container.add_stream("aac", rate=audio_rate, layout="mono") + frame_data = np.full((32, 32, 3), color, dtype=np.uint8) for _ in range(n_frames): frame = av.VideoFrame.from_ndarray(frame_data, format="rgb24") @@ -32,6 +46,19 @@ def _write_video(path, fps=24, seconds=1, color=(200, 100, 50)): container.mux(packet) for packet in stream.encode(): container.mux(packet) + + if audio_stream is not None: + n_samples = int(round(audio_rate * seconds)) + t = np.arange(n_samples, dtype=np.float32) / audio_rate + tone = (0.2 * np.sin(2 * np.pi * audio_hz * t)).reshape(1, n_samples) + aframe = av.AudioFrame.from_ndarray(np.ascontiguousarray(tone), format="fltp", layout="mono") + aframe.sample_rate = audio_rate + aframe.pts = 0 + for packet in audio_stream.encode(aframe): + container.mux(packet) + for packet in audio_stream.encode(): + container.mux(packet) + container.close() @@ -102,6 +129,26 @@ def test_force_rate_retimes_and_halves_frames(input_dir): assert float(components.frame_rate) == 12.0 +def test_audio_track_decoded_on_both_rate_paths(input_dir): + # force_rate=0 exercises _decode_audio_track's success path directly; force_rate=12.0 (half + # the clip's native 24fps) forces get_components() and proves audio_N still comes through — + # reused from components.audio rather than decoded a second time. + from nodes.multi_video_loader import _run_video + _write_video(input_dir / "clip.mp4", fps=24, seconds=1, audio_hz=440) + + out_cheap = _run_video(json.dumps([{"name": "clip.mp4"}]), force_rate=0.0, advanced=False) + audio_cheap = out_cheap[2] + assert isinstance(audio_cheap, dict) + assert audio_cheap["waveform"].numel() > 0 + assert audio_cheap["sample_rate"] > 0 + + out_forced = _run_video(json.dumps([{"name": "clip.mp4"}]), force_rate=12.0, advanced=False) + audio_forced = out_forced[2] + assert isinstance(audio_forced, dict) + assert audio_forced["waveform"].numel() > 0 + assert audio_forced["sample_rate"] > 0 + + def test_missing_file_raises(input_dir): from nodes.multi_video_loader import _run_video with pytest.raises(ValueError, match="gone.mp4"):