Load Media: decode only the frames a prepared video keeps, and only the first frame for a picture input

Memory now follows the prepared result, not the source file.

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01BzJLJWVu7RZnbRqm3nYAnX

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)

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"])

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)