feat: add VideoTrim and VideoCrop nodes with VIDEO_EDIT widget inputs

This commit is contained in:
Terry Jia
2026-08-23 15:44:13 -04:00
parent b78cec879b
commit 45d3ff0d87
8 changed files with 537 additions and 20 deletions

View File

@@ -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

View File

@@ -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)

View File

@@ -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",
]

View File

@@ -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",

View File

@@ -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