Skip to content

Commit 4a2058c

Browse files
committed
feat: add VideoTrim and VideoCrop nodes with VIDEO_EDIT widget inputs
1 parent 76135e5 commit 4a2058c

7 files changed

Lines changed: 415 additions & 11 deletions

File tree

comfy_api/latest/_input/video_types.py

Lines changed: 40 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
from typing import Optional, Union, IO
55
import io
66
import av
7-
from .._util import VideoContainer, VideoCodec, VideoComponents
7+
from .._util import VideoContainer, VideoCodec, VideoComponents, normalize_crop_rect
88

99
class VideoInput(ABC):
1010
"""
@@ -63,6 +63,45 @@ def as_trimmed(
6363
"""
6464
pass
6565

66+
def as_cropped(
67+
self,
68+
x: int = 0,
69+
y: int = 0,
70+
width: int = 0,
71+
height: int = 0,
72+
) -> VideoInput:
73+
"""
74+
Create a new VideoInput spatially cropped to the given pixel rectangle.
75+
76+
The rectangle is clamped to the frame and even-aligned for encoder
77+
compatibility. An empty or full-frame rectangle returns the input
78+
unchanged.
79+
80+
Default implementation materializes the video via get_components();
81+
subclasses should override with lazier strategies when possible.
82+
"""
83+
components = self.get_components()
84+
rect = normalize_crop_rect(
85+
x, y, width, height, components.images.shape[2], components.images.shape[1]
86+
)
87+
if rect is None:
88+
return self
89+
from .._input_impl.video_types import VideoFromComponents
90+
91+
cx, cy, cw, ch = rect
92+
return VideoFromComponents(
93+
VideoComponents(
94+
images=components.images[:, cy:cy + ch, cx:cx + cw, :],
95+
audio=components.audio,
96+
frame_rate=components.frame_rate,
97+
metadata=components.metadata,
98+
alpha=components.alpha[:, cy:cy + ch, cx:cx + cw]
99+
if components.alpha is not None
100+
else None,
101+
),
102+
bit_depth=self.get_bit_depth(),
103+
)
104+
66105
def get_stream_source(self) -> Union[str, io.BytesIO]:
67106
"""
68107
Get a streamable source for the video. This allows processing without

comfy_api/latest/_input_impl/video_types.py

Lines changed: 123 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,8 @@
1212
import math
1313
import os
1414
import torch
15-
from .._util import VideoContainer, VideoCodec, VideoComponents
15+
from .._util import VideoContainer, VideoCodec, VideoComponents, normalize_crop_rect
16+
import comfy.utils
1617
import logging
1718

1819

@@ -196,19 +197,25 @@ def webm_streams_compatible(streams) -> bool:
196197
return True
197198

198199

200+
def _rotation_quadrant(frame: av.VideoFrame) -> int:
201+
return int(round(frame.rotation // 90)) % 4 if frame.rotation else 0
202+
203+
199204
class VideoFromFile(VideoInput):
200205
"""
201206
Class representing video input from a file.
202207
"""
203208

204-
def __init__(self, file: str | io.BytesIO, *, start_time: float=0, duration: float=0):
209+
def __init__(self, file: str | io.BytesIO, *, start_time: float=0, duration: float=0,
210+
crop: tuple[int, int, int, int] | None = None):
205211
"""
206212
Initialize the VideoFromFile object based off of either a path on disk or a BytesIO object
207213
containing the file contents.
208214
"""
209215
self.__file = file
210216
self.__start_time = start_time
211217
self.__duration = duration
218+
self.__crop = crop
212219

213220
def get_stream_source(self) -> str | io.BytesIO:
214221
"""
@@ -238,7 +245,31 @@ def get_dimensions(self) -> tuple[int, int]:
238245
for stream in container.streams:
239246
if stream.type == 'video':
240247
assert isinstance(stream, av.VideoStream)
241-
return stream.width, stream.height
248+
if self.__crop is None:
249+
return stream.width, stream.height
250+
251+
display_width, display_height = self._get_display_dimensions()
252+
rect = normalize_crop_rect(*self.__crop, display_width, display_height)
253+
if rect is not None:
254+
return rect[2], rect[3]
255+
return display_width, display_height
256+
raise ValueError(f"No video stream found in file '{self.__file}'")
257+
258+
def _get_display_dimensions(self) -> tuple[int, int]:
259+
if isinstance(self.__file, io.BytesIO):
260+
self.__file.seek(0)
261+
with av.open(self.__file, mode='r') as container:
262+
for stream in container.streams:
263+
if stream.type == 'video':
264+
assert isinstance(stream, av.VideoStream)
265+
width, height = stream.width, stream.height
266+
try:
267+
frame = next(container.decode(stream), None)
268+
except av.error.FFmpegError:
269+
frame = None
270+
if frame is not None and _rotation_quadrant(frame) % 2:
271+
width, height = height, width
272+
return width, height
242273
raise ValueError(f"No video stream found in file '{self.__file}'")
243274

244275
def get_bit_depth(self) -> int:
@@ -415,6 +446,8 @@ def get_components_internal(self, container: InputContainer) -> VideoComponents:
415446
streams = [video_stream]
416447
has_first_audio_frame = False
417448
checked_alpha = False
449+
crop_rect = None
450+
crop_resolved = False
418451

419452
# Default to False so we decode until EOF if duration is 0
420453
video_done = False
@@ -485,9 +518,16 @@ def get_components_internal(self, container: InputContainer) -> VideoComponents:
485518
img = np.ascontiguousarray(align_graph[2].pull().to_ndarray(format=image_format)[:frame.height, :frame.width])
486519
else:
487520
img = frame.to_ndarray(format=image_format)
488-
if frame.rotation != 0:
489-
k = int(round(frame.rotation // 90))
490-
img = np.rot90(img, k=k, axes=(0, 1)).copy()
521+
rotation_quadrant = _rotation_quadrant(frame)
522+
if rotation_quadrant:
523+
img = np.rot90(img, k=rotation_quadrant, axes=(0, 1)).copy()
524+
if self.__crop is not None:
525+
if not crop_resolved:
526+
crop_rect = normalize_crop_rect(*self.__crop, img.shape[1], img.shape[0])
527+
crop_resolved = True
528+
if crop_rect is not None:
529+
cx, cy, cw, ch = crop_rect
530+
img = np.ascontiguousarray(img[cy:cy + ch, cx:cx + cw])
491531
if alphas is None:
492532
frames.append(torch.from_numpy(img))
493533
else:
@@ -586,6 +626,8 @@ def save_to(
586626
reuse_streams = False
587627
if self.__start_time or self.__duration:
588628
reuse_streams = False
629+
if self.__crop is not None:
630+
reuse_streams = False
589631

590632
if not reuse_streams:
591633
if bit_depth is None:
@@ -673,6 +715,16 @@ def _save_transcoded(
673715
if duration:
674716
duration_cap = math.ceil(duration * sample_rate)
675717

718+
if duration:
719+
window_seconds = duration
720+
else:
721+
try:
722+
window_seconds = max(self._get_raw_duration() - start_time, 0.0)
723+
except ValueError:
724+
window_seconds = 0.0
725+
progress_total = max(1, int(round(window_seconds * float(rate))))
726+
pbar = comfy.utils.ProgressBar(progress_total)
727+
676728
streams = [video_stream] if audio_stream is None else [video_stream, audio_stream]
677729
pts_step = max(1, int(round((1 / rate) / video_stream.time_base)))
678730
video_done = False
@@ -685,6 +737,8 @@ def _save_transcoded(
685737
source_size = None
686738
rotation_k = 0
687739
rotation_filter = None
740+
crop_rect = None
741+
crop_filter = None
688742
audio_started = False
689743
samples_written = 0
690744
pending_audio = []
@@ -758,13 +812,27 @@ def drain_audio(final=False):
758812
if end_pts is not None and frame.pts is not None:
759813
frame_duration = min(frame_duration, end_pts - frame.pts)
760814
if output is None:
761-
rotation_k = int(round(frame.rotation // 90)) % 4 if frame.rotation else 0
815+
rotation_k = _rotation_quadrant(frame)
762816
if rotation_k % 2:
763817
out_width, out_height = frame.height, frame.width
764818
else:
765819
out_width, out_height = frame.width, frame.height
820+
if self.__crop is not None:
821+
crop_rect = normalize_crop_rect(*self.__crop, out_width, out_height)
822+
if crop_rect is not None:
823+
out_width, out_height = crop_rect[2], crop_rect[3]
824+
if (out_width % 2 or out_height % 2) and crop_rect is None:
825+
even_width = out_width - out_width % 2
826+
even_height = out_height - out_height % 2
827+
if even_width > 0 and even_height > 0:
828+
crop_rect = (0, 0, even_width, even_height)
829+
out_width, out_height = even_width, even_height
766830
if out_width % 2 or out_height % 2:
767831
raise ValueError(f"{output_codec.value.upper()} output requires even dimensions, got {out_width}x{out_height}")
832+
if any(component.is_alpha for component in frame.format.components):
833+
logging.warning(
834+
"Transcoded video output does not support alpha; the alpha channel will be discarded."
835+
)
768836
source_size = (frame.width, frame.height)
769837
output = av.open(path, **open_kwargs)
770838
# Add metadata before writing any streams
@@ -810,6 +878,19 @@ def drain_audio(final=False):
810878
rotation_filter = (g_src, g_sink)
811879
rotation_filter[0].push(frame)
812880
frame = rotation_filter[1].pull()
881+
if crop_rect is not None:
882+
if crop_filter is None:
883+
g = av.filter.Graph()
884+
g_src = g.add_buffer(width=frame.width, height=frame.height,
885+
format=frame.format.name, time_base=video_stream.time_base)
886+
g_crop = g.add("crop", f"{crop_rect[2]}:{crop_rect[3]}:{crop_rect[0]}:{crop_rect[1]}")
887+
g_sink = g.add("buffersink")
888+
g_src.link_to(g_crop)
889+
g_crop.link_to(g_sink)
890+
g.configure()
891+
crop_filter = (g_src, g_sink)
892+
crop_filter[0].push(frame)
893+
frame = crop_filter[1].pull()
813894
if frame.color_range == ColorRange.JPEG and not preserve_source_color:
814895
# compress full-range sources (yuvj/MJPEG) to limited range
815896
frame = frame.reformat(format=pix_fmt, src_color_range="JPEG", dst_color_range="MPEG")
@@ -852,6 +933,7 @@ def drain_audio(final=False):
852933
out_packet.duration = video_frame_durations.pop(out_packet.pts, 0)
853934
output.mux(out_packet)
854935
drain_audio()
936+
pbar.update(1)
855937

856938
elif packet.stream == audio_stream and not audio_done:
857939
for resampled in itertools.chain.from_iterable(map(resampler.resample, packet.decode())):
@@ -921,11 +1003,42 @@ def as_trimmed(
9211003
self.get_stream_source(),
9221004
start_time=start_time + self.__start_time,
9231005
duration=duration,
1006+
crop=self.__crop,
9241007
)
925-
if trimmed.get_duration() < duration and strict_duration:
1008+
if strict_duration and duration and trimmed.get_duration() < duration:
9261009
return None
9271010
return trimmed
9281011

1012+
def as_cropped(
1013+
self, x: int = 0, y: int = 0, width: int = 0, height: int = 0
1014+
) -> VideoInput:
1015+
if int(width) <= 0 or int(height) <= 0:
1016+
return self
1017+
1018+
display_width, display_height = self._get_display_dimensions()
1019+
outer = (
1020+
normalize_crop_rect(*self.__crop, display_width, display_height)
1021+
if self.__crop is not None
1022+
else None
1023+
)
1024+
if outer is None:
1025+
rect = normalize_crop_rect(x, y, width, height, display_width, display_height)
1026+
else:
1027+
inner = normalize_crop_rect(x, y, width, height, outer[2], outer[3])
1028+
rect = (
1029+
(outer[0] + inner[0], outer[1] + inner[1], inner[2], inner[3])
1030+
if inner is not None
1031+
else None
1032+
)
1033+
if rect is None:
1034+
return self
1035+
return VideoFromFile(
1036+
self.get_stream_source(),
1037+
start_time=self.__start_time,
1038+
duration=self.__duration,
1039+
crop=rect,
1040+
)
1041+
9291042

9301043
class VideoFromComponents(VideoInput):
9311044
"""
@@ -942,6 +1055,8 @@ def get_components(self) -> VideoComponents:
9421055
images=self.__components.images,
9431056
audio=self.__components.audio,
9441057
frame_rate=self.__components.frame_rate,
1058+
metadata=self.__components.metadata,
1059+
alpha=self.__components.alpha,
9451060
)
9461061

9471062
def get_bit_depth(self) -> int:

comfy_api/latest/_io.py

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1416,6 +1416,41 @@ def __init__(self, id: str, display_name: str=None, optional=False, tooltip: str
14161416
self.default = []
14171417

14181418

1419+
@comfytype(io_type="VIDEO_EDIT")
1420+
class VideoEdit(ComfyTypeIO):
1421+
class VideoTrimSection(TypedDict):
1422+
start_time: float
1423+
duration: float
1424+
1425+
class VideoCropSection(TypedDict):
1426+
x: int
1427+
y: int
1428+
width: int
1429+
height: int
1430+
1431+
class VideoEditDict(TypedDict, total=False):
1432+
trim: 'VideoEdit.VideoTrimSection'
1433+
crop: 'VideoEdit.VideoCropSection'
1434+
Type = VideoEditDict
1435+
1436+
class Input(WidgetInput):
1437+
def __init__(self, id: str, display_name: str=None, optional=False, tooltip: str=None,
1438+
socketless: bool=True, default: dict=None, features: list[str]=None, advanced: bool=None):
1439+
super().__init__(id, display_name, optional, tooltip, None, default, socketless, None, None, None, None, advanced)
1440+
self.features = features if features is not None else ["trim", "crop"]
1441+
if default is None:
1442+
self.default = {}
1443+
if "trim" in self.features:
1444+
self.default["trim"] = {"start_time": 0.0, "duration": 0.0}
1445+
if "crop" in self.features:
1446+
self.default["crop"] = {"x": 0, "y": 0, "width": 0, "height": 0}
1447+
1448+
def as_dict(self):
1449+
return super().as_dict() | prune_dict({
1450+
"features": self.features,
1451+
})
1452+
1453+
14191454
@comfytype(io_type="HISTOGRAM")
14201455
class Histogram(ComfyTypeIO):
14211456
"""A histogram represented as a list of bin counts."""
@@ -2493,5 +2528,6 @@ def as_dict(self):
24932528
"Curve",
24942529
"Histogram",
24952530
"Range",
2531+
"VideoEdit",
24962532
"NodeReplace",
24972533
]

comfy_api/latest/_util/__init__.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from .video_types import VideoContainer, VideoCodec, VideoComponents
1+
from .video_types import VideoContainer, VideoCodec, VideoComponents, normalize_crop_rect
22
from .geometry_types import VOXEL, MESH, SPLAT, File3D
33
from .image_types import SVG
44

@@ -7,6 +7,7 @@
77
"VideoContainer",
88
"VideoCodec",
99
"VideoComponents",
10+
"normalize_crop_rect",
1011
"VOXEL",
1112
"MESH",
1213
"SPLAT",

comfy_api/latest/_util/video_types.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,3 +55,25 @@ class VideoComponents:
5555
audio: Optional[AudioInput] = None
5656
metadata: Optional[dict] = None
5757
alpha: Optional[MaskInput] = None
58+
59+
60+
def normalize_crop_rect(
61+
x: int, y: int, width: int, height: int, source_width: int, source_height: int
62+
) -> Optional[tuple[int, int, int, int]]:
63+
width = int(width)
64+
height = int(height)
65+
if width <= 0 or height <= 0:
66+
return None
67+
x = max(0, min(int(x), source_width - 1))
68+
y = max(0, min(int(y), source_height - 1))
69+
x -= x % 2
70+
y -= y % 2
71+
width = min(width, source_width - x)
72+
height = min(height, source_height - y)
73+
if x == 0 and y == 0 and width == source_width and height == source_height:
74+
return None
75+
width -= width % 2
76+
height -= height % 2
77+
if width <= 0 or height <= 0:
78+
return None
79+
return x, y, width, height

0 commit comments

Comments
 (0)