Claude
c70c74f01f
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BzJLJWVu7RZnbRqm3nYAnX
193 lines
10 KiB
Python
193 lines
10 KiB
Python
"""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. Keys: fps, skip_first, cut_last, every_nth, frame_cap, "
|
|
"width, height, fit (crop|pad|stretch), crop_left/right/top/bottom (percent), multiple_of. "
|
|
"Empty: all frames at 24 fps, size untouched."),
|
|
],
|
|
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)
|