diff --git a/modules/media/load/nodes.py b/modules/media/load/nodes.py index 26808e08..b819f3f4 100644 --- a/modules/media/load/nodes.py +++ b/modules/media/load/nodes.py @@ -46,6 +46,34 @@ def _read_audio(path): return {"waveform": torch.from_numpy(np.concatenate(chunks, axis=1)).unsqueeze(0), "sample_rate": rate} +def _read_prepared(path, options): + """A video's frames made ready, decoding only the ones that are kept and shaping each as it comes. + + Memory follows the prepared result, not the file: a long 4K clip with a frame cap and a target size never holds + more than the capped frames at the target size (plus the one being decoded). -> (frames, out_fps, start, span). + """ + import av + o = prep.merged(options) + with av.open(str(path)) as container: + stream = container.streams.video[0] + fps = float(stream.average_rate or REF_VIDEO_FPS) + count = sum(1 for packet in container.demux(stream) if packet.size) # counting packets decodes nothing + index, out_fps, start, span = prep.pick_frames(count, fps, o["fps"], o["skip_first"], o["cut_last"], o["every_nth"], o["frame_cap"]) + wanted, last, kept = set(index), max(index), {} + with av.open(str(path)) as container: + for i, frame in enumerate(container.decode(video=0)): + if i in wanted: + rgb = torch.from_numpy(frame.to_ndarray(format="rgb24")).float().div_(255)[None] + kept[i] = prep.shape(rgb, **prep.shape_options(o))[0] + if i >= last: + break + if not kept: + raise ValueError("no frames could be decoded") + have = sorted(kept) + pick = lambda i: kept[i] if i in kept else kept[max((h for h in have if h <= i), default=have[0])] # a short decode repeats the last frame + return torch.stack([pick(i) for i in index]), out_fps, start, span + + class FunUILoadMedia(io.ComfyNode): """Takes a picture, a video or an audio file and fills what it has: the app wires the output that fits.""" @@ -114,6 +142,31 @@ class FunUILoadMedia(io.ComfyNode): grey = torch.full((1, 64, 64, 3), 0.5) # no picture in it: the outputs for pictures are never wired return io.NodeOutput(grey, grey, sound) + wants_prepared = kind == "video" and full and (prepare or "").strip() + if wants_prepared: + try: + options = json.loads(prepare) + frames, out_fps, start, span = _read_prepared(path, options if isinstance(options, dict) else {}) + sound = _read_audio(path) + except (ValueError, TypeError) as exc: # bad JSON, or a cut that leaves nothing + raise RuntimeError(f"FunUI Load Media: {name!r} cannot be prepared ({exc}).") from exc + except Exception as exc: # noqa: BLE001 -- a damaged file raises whatever PyAV raises + raise RuntimeError(f"FunUI Load Media: could not read {name!r} ({exc}).") from exc + sound = prep.cut_audio(sound, start, span) if sound else _SILENCE + log.info("FunUI Load Media", f"loaded {name}: {frames.shape[0]} frames {frames.shape[2]}x{frames.shape[1]} at {out_fps:g} fps") + return io.NodeOutput(frames[:1], frames, sound) + + if kind == "video" and not full: + # A picture input wants the first frame: decoding the whole clip for it is what a long video cannot afford. + try: + import av + with av.open(str(path)) as container: + first = torch.from_numpy(next(container.decode(video=0)).to_ndarray(format="rgb24")).float().div_(255)[None] + except Exception as exc: # noqa: BLE001 -- a damaged file raises whatever PyAV raises + raise RuntimeError(f"FunUI Load Media: could not read {name!r} ({exc}).") from exc + log.info("FunUI Load Media", f"loaded {name}") + return io.NodeOutput(first, first, _SILENCE) + try: components = InputImpl.VideoFromFile(str(path)).get_components() except Exception as exc: # noqa: BLE001 -- a damaged file raises whatever PyAV raises @@ -127,15 +180,6 @@ class FunUILoadMedia(io.ComfyNode): return io.NodeOutput(first, first, _SILENCE) fps = float(components.frame_rate or REF_VIDEO_FPS) - if (prepare or "").strip(): - try: - options = json.loads(prepare) - frames, out_fps, start, span = prep.prepare(images, fps, options if isinstance(options, dict) else {}) - except (ValueError, TypeError) as exc: # bad JSON, or a cut that leaves nothing - raise RuntimeError(f"FunUI Load Media: {name!r} cannot be prepared ({exc}).") from exc - sound = prep.cut_audio(components.audio, start, span) if components.audio else _SILENCE - log.info("FunUI Load Media", f"loaded {name}: {frames.shape[0]} frames {frames.shape[2]}x{frames.shape[1]} at {out_fps:g} fps") - return io.NodeOutput(frames[:1], frames, sound) # Whatever the source rate, the model wants 24 fps: pick the frame showing at each 1/24 s. count = max(1, int(images.shape[0] / fps * REF_VIDEO_FPS)) index = torch.clamp((torch.arange(count) * fps / REF_VIDEO_FPS).long(), max=images.shape[0] - 1) diff --git a/modules/media/load/prepare.py b/modules/media/load/prepare.py index 687300a7..b691ca19 100644 --- a/modules/media/load/prepare.py +++ b/modules/media/load/prepare.py @@ -93,9 +93,20 @@ def shape(images, width=0, height=0, fit="crop", crop_left=0, crop_right=0, crop return out +def merged(options): + """The options with every missing key at its default; unknown keys dropped.""" + return {**DEFAULTS, **{k: v for k, v in (options or {}).items() if k in DEFAULTS}} + + +def shape_options(o): + """The part of the options that `shape` takes.""" + return dict(width=o["width"], height=o["height"], fit=o["fit"], crop_left=o["crop_left"], crop_right=o["crop_right"], + crop_top=o["crop_top"], crop_bottom=o["crop_bottom"], multiple_of=o["multiple_of"]) + + def prepare(images, src_fps, options=None): """The whole job. -> (frames, out_fps, start_sec, span_sec).""" - o = {**DEFAULTS, **{k: v for k, v in (options or {}).items() if k in DEFAULTS}} + o = merged(options) index, out_fps, start, span = pick_frames(images.shape[0], src_fps, o["fps"], o["skip_first"], o["cut_last"], o["every_nth"], o["frame_cap"]) frames = shape(images[torch.tensor(index)], o["width"], o["height"], o["fit"], o["crop_left"], o["crop_right"], o["crop_top"], o["crop_bottom"], o["multiple_of"]) diff --git a/modules/media/load/tests/test_load_media.py b/modules/media/load/tests/test_load_media.py index b5facfc4..1efb64b2 100644 --- a/modules/media/load/tests/test_load_media.py +++ b/modules/media/load/tests/test_load_media.py @@ -120,3 +120,47 @@ def test_slots_of_the_retired_video_and_audio_nodes_are_dropped_with_their_wires out = upgrade(slots, None) assert [s["id"] for s in out] == ["r2v", "ref"] assert out[0]["inputs"] == {"keep": 3, "b": ["ref", 0]} + + +def _long_clip(store, frames=30, rate=30, size=(64, 48)): + import av + import numpy as np + store.parent.mkdir(parents=True, exist_ok=True) + path = store.parent / "long.mp4" + with av.open(str(path), "w") as out: + stream = out.add_stream("libx264", rate=rate) + stream.width, stream.height, stream.pix_fmt = size[0], size[1], "yuv420p" + for i in range(frames): + arr = np.full((size[1], size[0], 3), (i * 7) % 255, dtype=np.uint8) + for packet in stream.encode(av.VideoFrame.from_ndarray(arr, format="rgb24")): + out.mux(packet) + for packet in stream.encode(): + out.mux(packet) + return media.save_upload("long.mp4", path.read_bytes())["id"] + + +def test_a_prepared_video_decodes_only_what_it_keeps_and_matches_the_whole_decode_route(store): + import json + import torch + from comfy_api.latest import InputImpl + from modules.media.load import prepare as prep + from modules.media.load.nodes import FunUILoadMedia + mid = _long_clip(store) + options = {"fps": 24, "skip_first": 2, "cut_last": 3, "every_nth": 2, "frame_cap": 6, "width": 32, "height": 32, "fit": "crop", "multiple_of": 16} + + _, frames, _ = FunUILoadMedia.execute(mid, full=True, prepare=json.dumps(options)).result + + comp = InputImpl.VideoFromFile(str(media.path_for(mid))).get_components() + expected, _, _, _ = prep.prepare(comp.images, float(comp.frame_rate), options) + assert frames.shape == expected.shape and frames.shape[0] == 6 + assert torch.allclose(frames, expected, atol=1 / 255) + + +def test_a_cut_that_leaves_nothing_is_explained_and_a_picture_input_reads_one_frame(store): + import json + from modules.media.load.nodes import FunUILoadMedia + mid = _long_clip(store) + with pytest.raises(RuntimeError, match="cannot be prepared"): + FunUILoadMedia.execute(mid, full=True, prepare=json.dumps({"skip_first": 100})).result + image, _, _ = FunUILoadMedia.execute(mid).result + assert image.shape == (1, 48, 64, 3)