diff --git a/comfy_api/latest/_input/video_types.py b/comfy_api/latest/_input/video_types.py index 2fbdbb341..600e116cc 100644 --- a/comfy_api/latest/_input/video_types.py +++ b/comfy_api/latest/_input/video_types.py @@ -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 diff --git a/comfy_api/latest/_input_impl/video_types.py b/comfy_api/latest/_input_impl/video_types.py index 34c6143ef..ca67c80ce 100644 --- a/comfy_api/latest/_input_impl/video_types.py +++ b/comfy_api/latest/_input_impl/video_types.py @@ -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) diff --git a/comfy_api/latest/_io.py b/comfy_api/latest/_io.py index 6d9e08f47..6c293b22d 100644 --- a/comfy_api/latest/_io.py +++ b/comfy_api/latest/_io.py @@ -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", ] diff --git a/comfy_api/latest/_util/__init__.py b/comfy_api/latest/_util/__init__.py index b27f5a97e..c3c9f0183 100644 --- a/comfy_api/latest/_util/__init__.py +++ b/comfy_api/latest/_util/__init__.py @@ -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", diff --git a/comfy_api/latest/_util/video_types.py b/comfy_api/latest/_util/video_types.py index 9eaaf0adc..3451034ac 100644 --- a/comfy_api/latest/_util/video_types.py +++ b/comfy_api/latest/_util/video_types.py @@ -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 diff --git a/comfy_extras/nodes_video.py b/comfy_extras/nodes_video.py index 58f58aaf4..846a94861 100644 --- a/comfy_extras/nodes_video.py +++ b/comfy_extras/nodes_video.py @@ -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: diff --git a/tests-unit/comfy_api_test/video_types_test.py b/tests-unit/comfy_api_test/video_types_test.py index fe7b8de2a..a24419ad9 100644 --- a/tests-unit/comfy_api_test/video_types_test.py +++ b/tests-unit/comfy_api_test/video_types_test.py @@ -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) diff --git a/tests-unit/comfy_extras_test/test_video_preview_nodes.py b/tests-unit/comfy_extras_test/test_video_preview_nodes.py new file mode 100644 index 000000000..b041cd578 --- /dev/null +++ b/tests-unit/comfy_extras_test/test_video_preview_nodes.py @@ -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