mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-08-25 02:20:18 +08:00
feat: add VideoTrim and VideoCrop nodes with VIDEO_EDIT widget inputs
This commit is contained in:
@@ -4,7 +4,7 @@ from fractions import Fraction
|
||||
from typing import Optional, Union, IO
|
||||
import io
|
||||
import av
|
||||
from .._util import VideoContainer, VideoCodec, VideoComponents
|
||||
from .._util import VideoContainer, VideoCodec, VideoComponents, normalize_crop_rect
|
||||
|
||||
class VideoInput(ABC):
|
||||
"""
|
||||
@@ -31,12 +31,15 @@ class VideoInput(ABC):
|
||||
bit_depth: int | None = None,
|
||||
crf: float | None = None,
|
||||
color_space: str | None = None,
|
||||
preset: str | None = None,
|
||||
):
|
||||
"""
|
||||
Abstract method to save the video input to a file.
|
||||
|
||||
bit_depth selects the encoded bit depth; None keeps the video's native depth.
|
||||
crf selects the H.264 or AV1 constant rate factor; None uses the encoder default.
|
||||
preset selects the H.264 encoder speed/compression trade-off (e.g. "ultrafast");
|
||||
None uses the encoder default. Ignored for other codecs.
|
||||
color_space="sRGB" selects SDR BT.709/sRGB, "HDR" selects BT.2020/HLG, and "HDR PQ"
|
||||
selects BT.2020/PQ. Bit depth is selected independently.
|
||||
Tensor-created videos default to sRGB when color_space is None. Loaded videos keep matching recognized native color
|
||||
@@ -63,6 +66,45 @@ class VideoInput(ABC):
|
||||
"""
|
||||
pass
|
||||
|
||||
def as_cropped(
|
||||
self,
|
||||
x: int = 0,
|
||||
y: int = 0,
|
||||
width: int = 0,
|
||||
height: int = 0,
|
||||
) -> VideoInput:
|
||||
"""
|
||||
Create a new VideoInput spatially cropped to the given pixel rectangle.
|
||||
|
||||
The rectangle is clamped to the frame and even-aligned for encoder
|
||||
compatibility. An empty or full-frame rectangle returns the input
|
||||
unchanged.
|
||||
|
||||
Default implementation materializes the video via get_components();
|
||||
subclasses should override with lazier strategies when possible.
|
||||
"""
|
||||
components = self.get_components()
|
||||
rect = normalize_crop_rect(
|
||||
x, y, width, height, components.images.shape[2], components.images.shape[1]
|
||||
)
|
||||
if rect is None:
|
||||
return self
|
||||
from .._input_impl.video_types import VideoFromComponents
|
||||
|
||||
cx, cy, cw, ch = rect
|
||||
return VideoFromComponents(
|
||||
VideoComponents(
|
||||
images=components.images[:, cy:cy + ch, cx:cx + cw, :].clone(),
|
||||
audio=components.audio,
|
||||
frame_rate=components.frame_rate,
|
||||
metadata=components.metadata,
|
||||
alpha=components.alpha[:, cy:cy + ch, cx:cx + cw].clone()
|
||||
if components.alpha is not None
|
||||
else None,
|
||||
),
|
||||
bit_depth=self.get_bit_depth(),
|
||||
)
|
||||
|
||||
def get_stream_source(self) -> Union[str, io.BytesIO]:
|
||||
"""
|
||||
Get a streamable source for the video. This allows processing without
|
||||
|
||||
@@ -12,7 +12,8 @@ import numpy as np
|
||||
import math
|
||||
import os
|
||||
import torch
|
||||
from .._util import VideoContainer, VideoCodec, VideoComponents
|
||||
from .._util import VideoContainer, VideoCodec, VideoComponents, normalize_crop_rect
|
||||
import comfy.utils
|
||||
import logging
|
||||
|
||||
|
||||
@@ -180,12 +181,18 @@ def video_stream_color_space(stream) -> str | None:
|
||||
return VIDEO_TRANSFER_COLOR_SPACES.get(stream.color_trc)
|
||||
|
||||
|
||||
def video_encoder_options(codec: VideoCodec, crf: float | None) -> dict[str, str]:
|
||||
if crf is None:
|
||||
return {}
|
||||
if codec == VideoCodec.AV1 and crf == 0:
|
||||
return {"svtav1-params": "lossless=1"}
|
||||
return {"crf": str(crf)}
|
||||
def video_encoder_options(
|
||||
codec: VideoCodec, crf: float | None, preset: str | None = None
|
||||
) -> dict[str, str]:
|
||||
options = {}
|
||||
if preset is not None and codec == VideoCodec.H264:
|
||||
options["preset"] = preset
|
||||
if crf is not None:
|
||||
if codec == VideoCodec.AV1 and crf == 0:
|
||||
options["svtav1-params"] = "lossless=1"
|
||||
else:
|
||||
options["crf"] = str(crf)
|
||||
return options
|
||||
|
||||
|
||||
def webm_streams_compatible(streams) -> bool:
|
||||
@@ -196,12 +203,17 @@ def webm_streams_compatible(streams) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
def _rotation_quadrant(frame: av.VideoFrame) -> int:
|
||||
return int(round(frame.rotation // 90)) % 4 if frame.rotation else 0
|
||||
|
||||
|
||||
class VideoFromFile(VideoInput):
|
||||
"""
|
||||
Class representing video input from a file.
|
||||
"""
|
||||
|
||||
def __init__(self, file: str | io.BytesIO, *, start_time: float=0, duration: float=0):
|
||||
def __init__(self, file: str | io.BytesIO, *, start_time: float=0, duration: float=0,
|
||||
crop: tuple[int, int, int, int] | None = None):
|
||||
"""
|
||||
Initialize the VideoFromFile object based off of either a path on disk or a BytesIO object
|
||||
containing the file contents.
|
||||
@@ -209,6 +221,7 @@ class VideoFromFile(VideoInput):
|
||||
self.__file = file
|
||||
self.__start_time = start_time
|
||||
self.__duration = duration
|
||||
self.__crop = crop
|
||||
|
||||
def get_stream_source(self) -> str | io.BytesIO:
|
||||
"""
|
||||
@@ -238,7 +251,31 @@ class VideoFromFile(VideoInput):
|
||||
for stream in container.streams:
|
||||
if stream.type == 'video':
|
||||
assert isinstance(stream, av.VideoStream)
|
||||
return stream.width, stream.height
|
||||
if self.__crop is None:
|
||||
return stream.width, stream.height
|
||||
|
||||
display_width, display_height = self._get_display_dimensions()
|
||||
rect = normalize_crop_rect(*self.__crop, display_width, display_height)
|
||||
if rect is not None:
|
||||
return rect[2], rect[3]
|
||||
return display_width, display_height
|
||||
raise ValueError(f"No video stream found in file '{self.__file}'")
|
||||
|
||||
def _get_display_dimensions(self) -> tuple[int, int]:
|
||||
if isinstance(self.__file, io.BytesIO):
|
||||
self.__file.seek(0)
|
||||
with av.open(self.__file, mode='r') as container:
|
||||
for stream in container.streams:
|
||||
if stream.type == 'video':
|
||||
assert isinstance(stream, av.VideoStream)
|
||||
width, height = stream.width, stream.height
|
||||
try:
|
||||
frame = next(container.decode(stream), None)
|
||||
except av.error.FFmpegError:
|
||||
frame = None
|
||||
if frame is not None and _rotation_quadrant(frame) % 2:
|
||||
width, height = height, width
|
||||
return width, height
|
||||
raise ValueError(f"No video stream found in file '{self.__file}'")
|
||||
|
||||
def get_bit_depth(self) -> int:
|
||||
@@ -395,6 +432,7 @@ class VideoFromFile(VideoInput):
|
||||
|
||||
def get_components_internal(self, container: InputContainer) -> VideoComponents:
|
||||
video_stream = self._get_first_video_stream(container)
|
||||
video_stream.thread_type = "AUTO"
|
||||
start_time, duration = self.get_active_trim_window()
|
||||
|
||||
# Get video frames
|
||||
@@ -415,6 +453,8 @@ class VideoFromFile(VideoInput):
|
||||
streams = [video_stream]
|
||||
has_first_audio_frame = False
|
||||
checked_alpha = False
|
||||
crop_rect = None
|
||||
crop_resolved = False
|
||||
|
||||
# Default to False so we decode until EOF if duration is 0
|
||||
video_done = False
|
||||
@@ -485,9 +525,16 @@ class VideoFromFile(VideoInput):
|
||||
img = np.ascontiguousarray(align_graph[2].pull().to_ndarray(format=image_format)[:frame.height, :frame.width])
|
||||
else:
|
||||
img = frame.to_ndarray(format=image_format)
|
||||
if frame.rotation != 0:
|
||||
k = int(round(frame.rotation // 90))
|
||||
img = np.rot90(img, k=k, axes=(0, 1)).copy()
|
||||
rotation_quadrant = _rotation_quadrant(frame)
|
||||
if rotation_quadrant:
|
||||
img = np.rot90(img, k=rotation_quadrant, axes=(0, 1)).copy()
|
||||
if self.__crop is not None:
|
||||
if not crop_resolved:
|
||||
crop_rect = normalize_crop_rect(*self.__crop, img.shape[1], img.shape[0])
|
||||
crop_resolved = True
|
||||
if crop_rect is not None:
|
||||
cx, cy, cw, ch = crop_rect
|
||||
img = np.ascontiguousarray(img[cy:cy + ch, cx:cx + cw])
|
||||
if alphas is None:
|
||||
frames.append(torch.from_numpy(img))
|
||||
else:
|
||||
@@ -554,6 +601,7 @@ class VideoFromFile(VideoInput):
|
||||
bit_depth: int | None = None,
|
||||
crf: float | None = None,
|
||||
color_space: str | None = None,
|
||||
preset: str | None = None,
|
||||
):
|
||||
if color_space is not None and color_space not in VIDEO_COLOR_TRANSFERS:
|
||||
raise ValueError(f"Unsupported video color space: {color_space}")
|
||||
@@ -586,11 +634,13 @@ class VideoFromFile(VideoInput):
|
||||
reuse_streams = False
|
||||
if self.__start_time or self.__duration:
|
||||
reuse_streams = False
|
||||
if self.__crop is not None:
|
||||
reuse_streams = False
|
||||
|
||||
if not reuse_streams:
|
||||
if bit_depth is None:
|
||||
bit_depth = source_bit_depth
|
||||
return self._save_transcoded(container, path, format=format, codec=codec, metadata=metadata, bit_depth=bit_depth, crf=crf, color_space=color_space)
|
||||
return self._save_transcoded(container, path, format=format, codec=codec, metadata=metadata, bit_depth=bit_depth, crf=crf, color_space=color_space, preset=preset)
|
||||
|
||||
streams = container.streams
|
||||
|
||||
@@ -625,10 +675,12 @@ class VideoFromFile(VideoInput):
|
||||
bit_depth: int,
|
||||
crf: float | None = None,
|
||||
color_space: str | None = None,
|
||||
preset: str | None = None,
|
||||
):
|
||||
"""Re-encode one frame at a time; peak memory does not scale with video length."""
|
||||
open_kwargs, output_format, output_codec = video_output_config(path, format, codec)
|
||||
video_stream = self._get_first_video_stream(container)
|
||||
video_stream.thread_type = "AUTO"
|
||||
start_time, duration = self.get_active_trim_window()
|
||||
start_pts = int(start_time / video_stream.time_base)
|
||||
end_pts = int((start_time + duration) / video_stream.time_base) if duration else None
|
||||
@@ -671,6 +723,16 @@ class VideoFromFile(VideoInput):
|
||||
if duration:
|
||||
duration_cap = math.ceil(duration * sample_rate)
|
||||
|
||||
if duration:
|
||||
window_seconds = duration
|
||||
else:
|
||||
try:
|
||||
window_seconds = max(self._get_raw_duration() - start_time, 0.0)
|
||||
except ValueError:
|
||||
window_seconds = 0.0
|
||||
progress_total = max(1, int(round(window_seconds * float(rate))))
|
||||
pbar = comfy.utils.ProgressBar(progress_total)
|
||||
|
||||
streams = [video_stream] if audio_stream is None else [video_stream, audio_stream]
|
||||
pts_step = max(1, int(round((1 / rate) / video_stream.time_base)))
|
||||
video_done = False
|
||||
@@ -683,6 +745,8 @@ class VideoFromFile(VideoInput):
|
||||
source_size = None
|
||||
rotation_k = 0
|
||||
rotation_filter = None
|
||||
crop_rect = None
|
||||
crop_filter = None
|
||||
audio_started = False
|
||||
samples_written = 0
|
||||
pending_audio = []
|
||||
@@ -756,13 +820,27 @@ class VideoFromFile(VideoInput):
|
||||
if end_pts is not None and frame.pts is not None:
|
||||
frame_duration = min(frame_duration, end_pts - frame.pts)
|
||||
if output is None:
|
||||
rotation_k = int(round(frame.rotation // 90)) % 4 if frame.rotation else 0
|
||||
rotation_k = _rotation_quadrant(frame)
|
||||
if rotation_k % 2:
|
||||
out_width, out_height = frame.height, frame.width
|
||||
else:
|
||||
out_width, out_height = frame.width, frame.height
|
||||
if self.__crop is not None:
|
||||
crop_rect = normalize_crop_rect(*self.__crop, out_width, out_height)
|
||||
if crop_rect is not None:
|
||||
out_width, out_height = crop_rect[2], crop_rect[3]
|
||||
if (out_width % 2 or out_height % 2) and crop_rect is None:
|
||||
even_width = out_width - out_width % 2
|
||||
even_height = out_height - out_height % 2
|
||||
if even_width > 0 and even_height > 0:
|
||||
crop_rect = (0, 0, even_width, even_height)
|
||||
out_width, out_height = even_width, even_height
|
||||
if out_width % 2 or out_height % 2:
|
||||
raise ValueError(f"{output_codec.value.upper()} output requires even dimensions, got {out_width}x{out_height}")
|
||||
if any(component.is_alpha for component in frame.format.components):
|
||||
logging.warning(
|
||||
"Transcoded video output does not support alpha; the alpha channel will be discarded."
|
||||
)
|
||||
source_size = (frame.width, frame.height)
|
||||
output = av.open(path, **open_kwargs)
|
||||
# Add metadata before writing any streams
|
||||
@@ -774,7 +852,7 @@ class VideoFromFile(VideoInput):
|
||||
out_video.width = out_width
|
||||
out_video.height = out_height
|
||||
out_video.pix_fmt = pix_fmt
|
||||
out_video.options = video_encoder_options(output_codec, crf)
|
||||
out_video.options = video_encoder_options(output_codec, crf, preset)
|
||||
if preserve_source_color:
|
||||
copy_color_properties(video_stream, out_video.codec_context)
|
||||
elif color_space is not None:
|
||||
@@ -808,6 +886,19 @@ class VideoFromFile(VideoInput):
|
||||
rotation_filter = (g_src, g_sink)
|
||||
rotation_filter[0].push(frame)
|
||||
frame = rotation_filter[1].pull()
|
||||
if crop_rect is not None:
|
||||
if crop_filter is None:
|
||||
g = av.filter.Graph()
|
||||
g_src = g.add_buffer(width=frame.width, height=frame.height,
|
||||
format=frame.format.name, time_base=video_stream.time_base)
|
||||
g_crop = g.add("crop", f"{crop_rect[2]}:{crop_rect[3]}:{crop_rect[0]}:{crop_rect[1]}")
|
||||
g_sink = g.add("buffersink")
|
||||
g_src.link_to(g_crop)
|
||||
g_crop.link_to(g_sink)
|
||||
g.configure()
|
||||
crop_filter = (g_src, g_sink)
|
||||
crop_filter[0].push(frame)
|
||||
frame = crop_filter[1].pull()
|
||||
if frame.color_range == ColorRange.JPEG and not preserve_source_color:
|
||||
# compress full-range sources (yuvj/MJPEG) to limited range
|
||||
frame = frame.reformat(format=pix_fmt, src_color_range="JPEG", dst_color_range="MPEG")
|
||||
@@ -850,6 +941,7 @@ class VideoFromFile(VideoInput):
|
||||
out_packet.duration = video_frame_durations.pop(out_packet.pts, 0)
|
||||
output.mux(out_packet)
|
||||
drain_audio()
|
||||
pbar.update(1)
|
||||
|
||||
elif packet.stream == audio_stream and not audio_done:
|
||||
for resampled in itertools.chain.from_iterable(map(resampler.resample, packet.decode())):
|
||||
@@ -919,11 +1011,42 @@ class VideoFromFile(VideoInput):
|
||||
self.get_stream_source(),
|
||||
start_time=start_time + self.__start_time,
|
||||
duration=duration,
|
||||
crop=self.__crop,
|
||||
)
|
||||
if trimmed.get_duration() < duration and strict_duration:
|
||||
if strict_duration and duration and trimmed.get_duration() < duration:
|
||||
return None
|
||||
return trimmed
|
||||
|
||||
def as_cropped(
|
||||
self, x: int = 0, y: int = 0, width: int = 0, height: int = 0
|
||||
) -> VideoInput:
|
||||
if int(width) <= 0 or int(height) <= 0:
|
||||
return self
|
||||
|
||||
display_width, display_height = self._get_display_dimensions()
|
||||
outer = (
|
||||
normalize_crop_rect(*self.__crop, display_width, display_height)
|
||||
if self.__crop is not None
|
||||
else None
|
||||
)
|
||||
if outer is None:
|
||||
rect = normalize_crop_rect(x, y, width, height, display_width, display_height)
|
||||
else:
|
||||
inner = normalize_crop_rect(x, y, width, height, outer[2], outer[3])
|
||||
rect = (
|
||||
(outer[0] + inner[0], outer[1] + inner[1], inner[2], inner[3])
|
||||
if inner is not None
|
||||
else None
|
||||
)
|
||||
if rect is None:
|
||||
return self
|
||||
return VideoFromFile(
|
||||
self.get_stream_source(),
|
||||
start_time=self.__start_time,
|
||||
duration=self.__duration,
|
||||
crop=rect,
|
||||
)
|
||||
|
||||
|
||||
class VideoFromComponents(VideoInput):
|
||||
"""
|
||||
@@ -943,6 +1066,8 @@ class VideoFromComponents(VideoInput):
|
||||
images=self.__components.images,
|
||||
audio=self.__components.audio,
|
||||
frame_rate=self.__components.frame_rate,
|
||||
metadata=self.__components.metadata,
|
||||
alpha=self.__components.alpha,
|
||||
)
|
||||
|
||||
def get_bit_depth(self) -> int:
|
||||
@@ -960,6 +1085,7 @@ class VideoFromComponents(VideoInput):
|
||||
bit_depth: int | None = None,
|
||||
crf: float | None = None,
|
||||
color_space: str | None = None,
|
||||
preset: str | None = None,
|
||||
):
|
||||
"""Save the video to a file path or BytesIO buffer."""
|
||||
if color_space is None:
|
||||
@@ -984,7 +1110,7 @@ class VideoFromComponents(VideoInput):
|
||||
video_stream.width = self.__components.images.shape[2]
|
||||
video_stream.height = self.__components.images.shape[1]
|
||||
video_stream.pix_fmt = pix_fmt
|
||||
video_stream.options = video_encoder_options(output_codec, crf)
|
||||
video_stream.options = video_encoder_options(output_codec, crf, preset)
|
||||
if color_space is not None:
|
||||
set_video_color_properties(video_stream.codec_context, color_space)
|
||||
|
||||
|
||||
@@ -1416,6 +1416,41 @@ class BoundingBoxes(ComfyTypeIO):
|
||||
self.default = []
|
||||
|
||||
|
||||
@comfytype(io_type="VIDEO_EDIT")
|
||||
class VideoEdit(ComfyTypeIO):
|
||||
class VideoTrimSection(TypedDict):
|
||||
start_time: float
|
||||
duration: float
|
||||
|
||||
class VideoCropSection(TypedDict):
|
||||
x: int
|
||||
y: int
|
||||
width: int
|
||||
height: int
|
||||
|
||||
class VideoEditDict(TypedDict, total=False):
|
||||
trim: 'VideoEdit.VideoTrimSection'
|
||||
crop: 'VideoEdit.VideoCropSection'
|
||||
Type = VideoEditDict
|
||||
|
||||
class Input(WidgetInput):
|
||||
def __init__(self, id: str, display_name: str=None, optional=False, tooltip: str=None,
|
||||
socketless: bool=True, default: dict=None, features: list[str]=None, advanced: bool=None):
|
||||
super().__init__(id, display_name, optional, tooltip, None, default, socketless, None, None, None, None, advanced)
|
||||
self.features = features if features is not None else ["trim", "crop"]
|
||||
if default is None:
|
||||
self.default = {}
|
||||
if "trim" in self.features:
|
||||
self.default["trim"] = {"start_time": 0.0, "duration": 0.0}
|
||||
if "crop" in self.features:
|
||||
self.default["crop"] = {"x": 0, "y": 0, "width": 0, "height": 0}
|
||||
|
||||
def as_dict(self):
|
||||
return super().as_dict() | prune_dict({
|
||||
"features": self.features,
|
||||
})
|
||||
|
||||
|
||||
@comfytype(io_type="HISTOGRAM")
|
||||
class Histogram(ComfyTypeIO):
|
||||
"""A histogram represented as a list of bin counts."""
|
||||
@@ -2493,5 +2528,6 @@ __all__ = [
|
||||
"Curve",
|
||||
"Histogram",
|
||||
"Range",
|
||||
"VideoEdit",
|
||||
"NodeReplace",
|
||||
]
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
from .video_types import VideoContainer, VideoCodec, VideoComponents
|
||||
from .video_types import VideoContainer, VideoCodec, VideoComponents, normalize_crop_rect
|
||||
from .geometry_types import VOXEL, MESH, SPLAT, File3D
|
||||
from .image_types import SVG
|
||||
|
||||
@@ -7,6 +7,7 @@ __all__ = [
|
||||
"VideoContainer",
|
||||
"VideoCodec",
|
||||
"VideoComponents",
|
||||
"normalize_crop_rect",
|
||||
"VOXEL",
|
||||
"MESH",
|
||||
"SPLAT",
|
||||
|
||||
@@ -55,3 +55,25 @@ class VideoComponents:
|
||||
audio: Optional[AudioInput] = None
|
||||
metadata: Optional[dict] = None
|
||||
alpha: Optional[MaskInput] = None
|
||||
|
||||
|
||||
def normalize_crop_rect(
|
||||
x: int, y: int, width: int, height: int, source_width: int, source_height: int
|
||||
) -> Optional[tuple[int, int, int, int]]:
|
||||
width = int(width)
|
||||
height = int(height)
|
||||
if width <= 0 or height <= 0:
|
||||
return None
|
||||
x = max(0, min(int(x), source_width - 1))
|
||||
y = max(0, min(int(y), source_height - 1))
|
||||
x -= x % 2
|
||||
y -= y % 2
|
||||
width = min(width, source_width - x)
|
||||
height = min(height, source_height - y)
|
||||
if x == 0 and y == 0 and width == source_width and height == source_height:
|
||||
return None
|
||||
width -= width % 2
|
||||
height -= height % 2
|
||||
if width <= 0 or height <= 0:
|
||||
return None
|
||||
return x, y, width, height
|
||||
|
||||
@@ -3,6 +3,7 @@ import av
|
||||
import torch
|
||||
import folder_paths
|
||||
import json
|
||||
import weakref
|
||||
from typing import Optional
|
||||
from typing_extensions import override
|
||||
from fractions import Fraction
|
||||
@@ -295,6 +296,7 @@ class LoadVideo(io.ComfyNode):
|
||||
display_name="Load Video",
|
||||
category="video",
|
||||
essentials_category="Basics",
|
||||
has_intermediate_output=True,
|
||||
inputs=[
|
||||
io.Combo.Input("file", options=sorted(files), upload=io.UploadType.video),
|
||||
],
|
||||
@@ -306,7 +308,8 @@ class LoadVideo(io.ComfyNode):
|
||||
@classmethod
|
||||
def execute(cls, file) -> io.NodeOutput:
|
||||
video_path = folder_paths.get_annotated_filepath(file)
|
||||
return io.NodeOutput(InputImpl.VideoFromFile(video_path))
|
||||
source = InputImpl.VideoFromFile(video_path)
|
||||
return io.NodeOutput(source, ui=preview_input_video(file, source))
|
||||
|
||||
@classmethod
|
||||
def fingerprint_inputs(s, file):
|
||||
@@ -323,6 +326,67 @@ class LoadVideo(io.ComfyNode):
|
||||
|
||||
return True
|
||||
|
||||
_preview_results: "weakref.WeakKeyDictionary[Input.Video, tuple[str, ui.SavedResult]]" = weakref.WeakKeyDictionary()
|
||||
|
||||
|
||||
def preview_input_video(file: str, video: Input.Video | None = None) -> ui.PreviewVideo:
|
||||
name, _ = folder_paths.annotated_filepath(file)
|
||||
subfolder, _, filename = name.replace("\\", "/").rpartition("/")
|
||||
result = ui.SavedResult(filename, subfolder, io.FolderType.input)
|
||||
if video is not None:
|
||||
_preview_results[video] = (folder_paths.get_annotated_filepath(file), result)
|
||||
return ui.PreviewVideo([result])
|
||||
|
||||
|
||||
def save_video_preview(video: Input.Video) -> ui.PreviewVideo:
|
||||
cached = _preview_results.get(video)
|
||||
if cached is not None and os.path.isfile(cached[0]):
|
||||
return ui.PreviewVideo([cached[1]])
|
||||
|
||||
full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(
|
||||
"ComfyUI_temp_video", folder_paths.get_temp_directory(), 0, 0
|
||||
)
|
||||
preview_format = Types.VideoContainer.MP4
|
||||
file = f"{filename}_{counter:05}_.{Types.VideoContainer.get_extension(preview_format)}"
|
||||
full_path = os.path.join(full_output_folder, file)
|
||||
video.save_to(
|
||||
full_path,
|
||||
format=preview_format,
|
||||
codec="auto",
|
||||
preset="ultrafast",
|
||||
)
|
||||
result = ui.SavedResult(file, subfolder, io.FolderType.temp)
|
||||
_preview_results[video] = (full_path, result)
|
||||
return ui.PreviewVideo([result])
|
||||
|
||||
|
||||
def apply_video_trim(video: Input.Video, trim, strict_duration: bool = False) -> Input.Video:
|
||||
trim = trim or {}
|
||||
start_time = float(trim.get("start_time", 0.0))
|
||||
duration = float(trim.get("duration", 0.0))
|
||||
if duration < 0:
|
||||
raise ValueError(f"Trim duration must be >= 0, got {duration}")
|
||||
if start_time == 0.0 and duration == 0.0:
|
||||
return video
|
||||
|
||||
trimmed = video.as_trimmed(start_time, duration, strict_duration=strict_duration)
|
||||
if trimmed is None:
|
||||
raise ValueError(
|
||||
f"Failed to trim video:\nSource duration: {video.get_duration()}\nStart time: {start_time}\nTarget duration: {duration}"
|
||||
)
|
||||
return trimmed
|
||||
|
||||
|
||||
def apply_video_crop(video: Input.Video, crop) -> Input.Video:
|
||||
crop = crop or {}
|
||||
return video.as_cropped(
|
||||
int(crop.get("x", 0)),
|
||||
int(crop.get("y", 0)),
|
||||
int(crop.get("width", 0)),
|
||||
int(crop.get("height", 0)),
|
||||
)
|
||||
|
||||
|
||||
class VideoSlice(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
@@ -370,6 +434,74 @@ class VideoSlice(io.ComfyNode):
|
||||
)
|
||||
|
||||
|
||||
class VideoTrim(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="VideoTrim",
|
||||
display_name="Trim Video (Advanced)",
|
||||
search_aliases=["trim video duration", "skip first frames", "cut video", "start time"],
|
||||
category="video",
|
||||
is_experimental=True,
|
||||
is_output_node=True,
|
||||
essentials_category="Video Tools",
|
||||
has_intermediate_output=True,
|
||||
inputs=[
|
||||
io.Video.Input("video"),
|
||||
io.VideoEdit.Input(
|
||||
"trim",
|
||||
features=["trim"],
|
||||
tooltip="Trim window in seconds. Duration 0 keeps the video until the end.",
|
||||
),
|
||||
io.Boolean.Input(
|
||||
"strict_duration",
|
||||
default=False,
|
||||
advanced=True,
|
||||
tooltip="If True, when the specified duration is not possible, an error will be raised.",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
io.Video.Output(),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, video: io.Video.Type, trim: io.VideoEdit.Type, strict_duration: bool) -> io.NodeOutput:
|
||||
trimmed = apply_video_trim(video, (trim or {}).get("trim"), strict_duration=strict_duration)
|
||||
return io.NodeOutput(trimmed, ui=save_video_preview(trimmed))
|
||||
|
||||
|
||||
class VideoCrop(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="VideoCrop",
|
||||
display_name="Crop Video",
|
||||
search_aliases=["crop video", "cut region", "spatial crop"],
|
||||
category="video",
|
||||
is_experimental=True,
|
||||
is_output_node=True,
|
||||
essentials_category="Video Tools",
|
||||
has_intermediate_output=True,
|
||||
inputs=[
|
||||
io.Video.Input("video"),
|
||||
io.VideoEdit.Input(
|
||||
"crop",
|
||||
features=["crop"],
|
||||
tooltip="Crop region in pixels. Zero width/height keeps the full frame.",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
io.Video.Output(),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, video: io.Video.Type, crop: io.VideoEdit.Type) -> io.NodeOutput:
|
||||
cropped = apply_video_crop(video, (crop or {}).get("crop"))
|
||||
return io.NodeOutput(cropped, ui=save_video_preview(cropped))
|
||||
|
||||
|
||||
class VideoExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||
@@ -380,6 +512,8 @@ class VideoExtension(ComfyExtension):
|
||||
GetVideoComponents,
|
||||
LoadVideo,
|
||||
VideoSlice,
|
||||
VideoTrim,
|
||||
VideoCrop,
|
||||
]
|
||||
|
||||
async def comfy_entrypoint() -> VideoExtension:
|
||||
|
||||
@@ -8,6 +8,7 @@ import io
|
||||
import numpy as np
|
||||
from fractions import Fraction
|
||||
from comfy_api.input_impl.video_types import VideoFromFile, VideoFromComponents
|
||||
from comfy_api.latest._util.video_types import normalize_crop_rect
|
||||
from comfy_api.util.video_types import VideoComponents, VideoContainer, VideoCodec
|
||||
from comfy_api.input.basic_types import AudioInput
|
||||
from av.error import InvalidDataError
|
||||
@@ -1304,3 +1305,106 @@ def test_save_to_transcode_skips_undecodable_audio():
|
||||
for path in (mixed, all_bad):
|
||||
if path:
|
||||
os.unlink(path)
|
||||
|
||||
|
||||
def test_as_trimmed_strict_duration_gates_unavailable_length(simple_video_file):
|
||||
video = VideoFromFile(simple_video_file)
|
||||
|
||||
assert video.as_trimmed(0.0, 10.0, strict_duration=True) is None
|
||||
|
||||
relaxed = video.as_trimmed(0.0, 10.0, strict_duration=False)
|
||||
assert relaxed is not None
|
||||
assert relaxed.get_duration() < 10.0
|
||||
assert relaxed.get_duration() == pytest.approx(video.get_duration(), abs=EPSILON)
|
||||
|
||||
|
||||
def test_normalize_crop_rect_aligns_odd_origin_to_chroma_grid():
|
||||
assert normalize_crop_rect(1, 1, 100, 100, 1920, 1080) == (0, 0, 100, 100)
|
||||
assert normalize_crop_rect(3, 5, 10, 9, 64, 48) == (2, 4, 10, 8)
|
||||
assert normalize_crop_rect(0, 0, 64, 48, 64, 48) is None
|
||||
|
||||
|
||||
def _create_marker_video(path, width=64, height=48, marker_x=2, frames=3, fps=8):
|
||||
with av.open(path, mode="w") as container:
|
||||
stream = container.add_stream("h264", rate=fps)
|
||||
stream.width = width
|
||||
stream.height = height
|
||||
stream.pix_fmt = "yuv420p"
|
||||
pixels = torch.zeros(height, width, 3, dtype=torch.uint8).numpy()
|
||||
pixels[:, marker_x:marker_x + 2, :] = 255
|
||||
for _ in range(frames):
|
||||
frame = av.VideoFrame.from_ndarray(pixels, format="rgb24")
|
||||
frame = frame.reformat(format="yuv420p")
|
||||
container.mux(stream.encode(frame))
|
||||
container.mux(stream.encode(None))
|
||||
|
||||
|
||||
def _brightest_column(images):
|
||||
return images[0].float().mean(dim=(0, 2)).argmax().item()
|
||||
|
||||
|
||||
def test_cropped_decode_and_save_paths_select_same_pixels(tmp_path):
|
||||
source = str(tmp_path / "marker.mp4")
|
||||
_create_marker_video(source, marker_x=4)
|
||||
|
||||
cropped = VideoFromFile(source).as_cropped(3, 1, 16, 16)
|
||||
|
||||
components = cropped.get_components()
|
||||
assert tuple(components.images.shape[1:3]) == (16, 16)
|
||||
decode_column = _brightest_column(components.images)
|
||||
assert decode_column in (2, 3)
|
||||
|
||||
saved = str(tmp_path / "cropped.mp4")
|
||||
cropped.save_to(saved)
|
||||
saved_components = VideoFromFile(saved).get_components()
|
||||
assert tuple(saved_components.images.shape[1:3]) == (16, 16)
|
||||
save_column = _brightest_column(saved_components.images)
|
||||
assert save_column in (2, 3)
|
||||
|
||||
assert decode_column == save_column
|
||||
|
||||
|
||||
def test_as_cropped_components_releases_uncropped_storage():
|
||||
images = torch.rand(2, 8, 8, 3)
|
||||
video = VideoFromComponents(
|
||||
VideoComponents(images=images, frame_rate=Fraction(8))
|
||||
)
|
||||
|
||||
cropped = video.as_cropped(0, 0, 4, 4)
|
||||
cropped_images = cropped.get_components().images
|
||||
|
||||
assert tuple(cropped_images.shape[1:3]) == (4, 4)
|
||||
assert (
|
||||
cropped_images.untyped_storage().data_ptr()
|
||||
!= images.untyped_storage().data_ptr()
|
||||
)
|
||||
|
||||
|
||||
def test_video_encoder_options_applies_h264_preset():
|
||||
from comfy_api.latest._input_impl.video_types import video_encoder_options
|
||||
|
||||
assert video_encoder_options(VideoCodec.H264, None, "ultrafast") == {
|
||||
"preset": "ultrafast"
|
||||
}
|
||||
assert video_encoder_options(VideoCodec.H264, 23.0, "ultrafast") == {
|
||||
"preset": "ultrafast",
|
||||
"crf": "23.0",
|
||||
}
|
||||
assert video_encoder_options(VideoCodec.H264, 23.0, None) == {"crf": "23.0"}
|
||||
assert video_encoder_options(VideoCodec.AV1, None, "ultrafast") == {}
|
||||
assert video_encoder_options(VideoCodec.AV1, 0, "ultrafast") == {
|
||||
"svtav1-params": "lossless=1"
|
||||
}
|
||||
|
||||
|
||||
def test_save_to_preset_transcodes_playable_output(tmp_path):
|
||||
source = create_test_video(width=32, height=32)
|
||||
try:
|
||||
out = str(tmp_path / "preset.mp4")
|
||||
VideoFromFile(source).as_cropped(0, 0, 16, 16).save_to(
|
||||
out, preset="ultrafast"
|
||||
)
|
||||
saved = VideoFromFile(out).get_components()
|
||||
assert tuple(saved.images.shape[1:3]) == (16, 16)
|
||||
finally:
|
||||
os.unlink(source)
|
||||
|
||||
52
tests-unit/comfy_extras_test/test_video_preview_nodes.py
Normal file
52
tests-unit/comfy_extras_test/test_video_preview_nodes.py
Normal file
@@ -0,0 +1,52 @@
|
||||
"""Tests for video preview generation in comfy_extras/nodes_video.py."""
|
||||
import os
|
||||
|
||||
import av
|
||||
import torch
|
||||
|
||||
import folder_paths
|
||||
from comfy_api.input_impl.video_types import VideoFromFile
|
||||
from comfy_extras import nodes_video
|
||||
|
||||
|
||||
def _make_video(path, width=64, height=48, frames=3, fps=8):
|
||||
with av.open(str(path), "w") as container:
|
||||
stream = container.add_stream("h264", rate=fps)
|
||||
stream.width = width
|
||||
stream.height = height
|
||||
stream.pix_fmt = "yuv420p"
|
||||
for _ in range(frames):
|
||||
frame = av.VideoFrame.from_ndarray(
|
||||
torch.zeros(height, width, 3, dtype=torch.uint8).numpy(),
|
||||
format="rgb24",
|
||||
)
|
||||
for packet in stream.encode(frame.reformat(format="yuv420p")):
|
||||
container.mux(packet)
|
||||
for packet in stream.encode(None):
|
||||
container.mux(packet)
|
||||
return str(path)
|
||||
|
||||
|
||||
def test_save_video_preview_encodes_cropped_video(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(folder_paths, "get_temp_directory", lambda: str(tmp_path))
|
||||
|
||||
source = _make_video(tmp_path / "src.mp4")
|
||||
cropped = VideoFromFile(source).as_cropped(0, 0, 32, 24)
|
||||
|
||||
preview = nodes_video.save_video_preview(cropped)
|
||||
entry = preview.as_dict()["images"][0]
|
||||
|
||||
preview_path = os.path.join(str(tmp_path), entry["subfolder"], entry["filename"])
|
||||
assert VideoFromFile(preview_path).get_dimensions() == (32, 24)
|
||||
|
||||
|
||||
def test_save_video_preview_reuses_cached_result(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(folder_paths, "get_temp_directory", lambda: str(tmp_path))
|
||||
|
||||
source = _make_video(tmp_path / "src.mp4")
|
||||
cropped = VideoFromFile(source).as_cropped(0, 0, 32, 24)
|
||||
|
||||
first = nodes_video.save_video_preview(cropped).as_dict()["images"][0]
|
||||
second = nodes_video.save_video_preview(cropped).as_dict()["images"][0]
|
||||
|
||||
assert second == first
|
||||
Reference in New Issue
Block a user