olivv-cs-ComfyUI-FunPack/modules/media/load/nodes.py

"""Load one file from FunUI's media library (core/media.py): a picture, a video or sound.

The bridge between "picked in the Media bin" and "wired into the graph": the
media store has no idea what a reference image or an aspect ratio source is, and
the graph has no idea what a media id is. This is the one place both are true.

`media_id` is a plain string widget rather than a picker combo on purpose --
the id comes from the app (the Media bin hands it to this node's input before
a run queues, the same way a scene's prompt text lands on CLIPTextEncode's
`text`), not from someone typing a file name into this node by hand.
"""

import json

import torch
from comfy_api.latest import InputImpl, io

from ..._core import log, media as media_mod
from . import prepare as prep


#: Frame rate the H3 reference-video input takes.
REF_VIDEO_FPS = 24
#: What the audio output holds when the file has no sound (or nothing was asked for): one silent sample.
_SILENCE = {"waveform": torch.zeros((1, 1, 1)), "sample_rate": 44100}


def _read_audio(path):
    """A file's sound as ComfyUI AUDIO, or None when it has none."""
    import av
    import numpy as np
    chunks, rate = [], 0
    with av.open(str(path)) as container:
        if not container.streams.audio:
            return None
        for frame in container.decode(audio=0):
            rate = frame.sample_rate
            data = frame.to_ndarray()
            if data.dtype.kind in "iu":
                data = data.astype(np.float32) / float(np.iinfo(data.dtype).max)
            if frame.format.is_packed:
                data = data.reshape(-1, len(frame.layout.channels)).T
            chunks.append(data.astype(np.float32))
    if not chunks:
        return None
    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."""

    @classmethod
    def define_schema(cls) -> io.Schema:
        return io.Schema(
            node_id="FunUILoadMedia",
            display_name="FunUI Load Media",
            category="FunUI/Media",
            description="Load a picture, a video or an audio file from FunUI's media library by id. "
                        "image: the picture (a video's first frame). frames: a video's frames at 24 fps. "
                        "audio: a video's soundtrack or an audio file.",
            inputs=[
                io.String.Input(
                    "media_id", default="",
                    tooltip="A media library id, picked in the Media bin."),
                io.Int.Input(
                    "blank_width", default=0, min=0, max=16384, optional=True,
                    tooltip="With nothing picked, make a plain grey picture this wide instead of refusing. 0 = refuse."),
                io.Int.Input(
                    "blank_height", default=0, min=0, max=16384, optional=True,
                    tooltip="With nothing picked, make a plain grey picture this tall instead of refusing. 0 = refuse."),
                io.Int.Input(
                    "blank_frames", default=0, min=0, max=100000, optional=True,
                    tooltip="With nothing picked and a size given: how many grey frames the stand-in video has. 0 = one."),
                io.Boolean.Input(
                    "full", default=False, optional=True,
                    tooltip="Read the whole video and its sound. Off: only the first frame, which is all a picture input needs."),
                io.String.Input(
                    "prepare", default="", optional=True,
                    tooltip="With full on: how to get the video ready, as JSON (fps, skip_first, cut_last, every_nth, frame_cap, width, height, fit, crop_*, multiple_of). Empty keeps all frames at 24 fps."),
            ],
            outputs=[io.Image.Output(display_name="image"), io.Image.Output(display_name="frames"),
                     io.Audio.Output(display_name="audio")],
        )

    @classmethod
    def execute(cls, media_id: str, blank_width: int = 0, blank_height: int = 0, blank_frames: int = 0, full: bool = False, prepare: str = "") -> io.NodeOutput:
        mid = (media_id or "").strip()
        if not mid and blank_width > 0 and blank_height > 0:
            # A text-only start: the picture exists only to give the pipeline a SHAPE (see the H3 preset).
            log.info("FunUI Load Media", f"nothing picked: a plain {blank_width}x{blank_height} picture stands in")
            grey = torch.full((1, blank_height, blank_width, 3), 0.5)
            frames = grey.expand(max(1, blank_frames), -1, -1, -1).contiguous() if blank_frames > 1 else grey
            return io.NodeOutput(grey, frames, _SILENCE)
        if not mid:
            raise RuntimeError(
                "FunUI Load Media: nothing is picked for this input. "
                "Choose a file in the Media bin.")
        entry = media_mod.get(mid)
        path = media_mod.path_for(mid)
        if entry is None or path is None:
            raise RuntimeError(
                f"FunUI Load Media: media {mid!r} is not in the library any "
                f"more -- it may have been deleted. Pick another.")
        name, kind = entry.get("name", mid), entry.get("kind")
        if kind not in ("image", "video", "audio"):
            raise RuntimeError(f"FunUI Load Media: {name!r} is {kind or 'a file'}, which is not a picture, a video or sound.")

        if kind == "audio":
            try:
                sound = _read_audio(path)
            except Exception:                                # noqa: BLE001 -- a damaged file raises whatever PyAV raises
                sound = None
            if sound is None:
                raise RuntimeError(f"FunUI Load Media: could not read {name!r}.")
            log.info("FunUI Load Media", f"loaded {name} (audio)")
            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
            raise RuntimeError(f"FunUI Load Media: could not read {name!r} ({exc}).") from exc
        images = components.images
        if images is None or images.shape[0] == 0:
            raise RuntimeError(f"FunUI Load Media: could not read {name!r}.")
        first = images[:1]
        if kind == "image" or not full:
            log.info("FunUI Load Media", f"loaded {name}")
            return io.NodeOutput(first, first, _SILENCE)

        fps = float(components.frame_rate or REF_VIDEO_FPS)
        # 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)
        sound = components.audio or _SILENCE
        log.info("FunUI Load Media", f"loaded {name}: {count} frames at {REF_VIDEO_FPS} fps")
        return io.NodeOutput(first, images[index], sound)