From d364d3f8b5cc2641214d1c2984b0b97b12708e82 Mon Sep 17 00:00:00 2001 From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com> Date: Tue, 14 Apr 2026 14:39:56 +0800 Subject: [PATCH 1/4] Init implementation of video dataset nodes --- comfy_extras/nodes_dataset.py | 476 +++++++++++++++++++++++++++++++++- 1 file changed, 473 insertions(+), 3 deletions(-) diff --git a/comfy_extras/nodes_dataset.py b/comfy_extras/nodes_dataset.py index 98ed25d7e..089d37dbc 100644 --- a/comfy_extras/nodes_dataset.py +++ b/comfy_extras/nodes_dataset.py @@ -2,6 +2,7 @@ import logging import os import json +import av import numpy as np import torch from PIL import Image @@ -42,6 +43,49 @@ def load_and_process_images(image_files, input_dir): return output_images +VALID_VIDEO_EXTENSIONS = [".mp4", ".avi", ".mov", ".webm", ".mkv", ".flv"] + + +def load_video_frames(video_path, max_frames=0, frame_stride=1, start_frame=0): + """Load video file and return frames as a tensor. + + Args: + video_path: Path to the video file + max_frames: Maximum number of frames to load (0 = all) + frame_stride: Sample every Nth frame + start_frame: Frame index to start from + + Returns: + torch.Tensor: Video frames as [T, H, W, C] float32 tensor in [0, 1] + """ + container = av.open(video_path) + stream = container.streams.video[0] + + frames = [] + frame_idx = 0 + for frame in container.decode(stream): + if frame_idx < start_frame: + frame_idx += 1 + continue + if (frame_idx - start_frame) % frame_stride != 0: + frame_idx += 1 + continue + if max_frames > 0 and len(frames) >= max_frames: + break + + img = frame.to_ndarray(format='rgb24') + img_tensor = torch.from_numpy(img.copy()).float() / 255.0 + frames.append(img_tensor) + frame_idx += 1 + + container.close() + + if not frames: + raise ValueError(f"No frames could be extracted from {video_path}") + + return torch.stack(frames) # [T, H, W, C] + + class LoadImageDataSetFromFolderNode(io.ComfyNode): @classmethod def define_schema(cls): @@ -153,6 +197,166 @@ class LoadImageTextDataSetFromFolderNode(io.ComfyNode): return io.NodeOutput(output_tensor, captions) +class LoadVideoDataSetFromFolderNode(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LoadVideoDataSetFromFolder", + display_name="Load Video Dataset from Folder", + category="dataset", + is_experimental=True, + inputs=[ + io.Combo.Input( + "folder", + options=folder_paths.get_input_subfolders(), + tooltip="The folder containing video files.", + ), + io.Int.Input( + "max_frames", + default=0, + min=0, + max=99999, + tooltip="Maximum frames to load per video (0 = all frames).", + ), + io.Int.Input( + "frame_stride", + default=1, + min=1, + max=1000, + tooltip="Sample every Nth frame (1 = every frame).", + ), + io.Int.Input( + "start_frame", + default=0, + min=0, + max=99999, + tooltip="Frame index to start loading from.", + ), + ], + outputs=[ + io.Image.Output( + display_name="videos", + is_output_list=True, + tooltip="List of video tensors, each [T, H, W, C].", + ), + ], + ) + + @classmethod + def execute(cls, folder, max_frames, frame_stride, start_frame): + sub_input_dir = os.path.join(folder_paths.get_input_directory(), folder) + video_files = sorted([ + f for f in os.listdir(sub_input_dir) + if any(f.lower().endswith(ext) for ext in VALID_VIDEO_EXTENSIONS) + ]) + + if not video_files: + raise ValueError(f"No video files found in {sub_input_dir}") + + output_videos = [] + for file in video_files: + video_path = os.path.join(sub_input_dir, file) + frames = load_video_frames(video_path, max_frames, frame_stride, start_frame) + output_videos.append(frames) + logging.info(f"Loaded {file}: {frames.shape[0]} frames, {frames.shape[1]}x{frames.shape[2]}") + + logging.info(f"Loaded {len(output_videos)} videos from {sub_input_dir}") + return io.NodeOutput(output_videos) + + +class LoadVideoTextDataSetFromFolderNode(io.ComfyNode): + @classmethod + def define_schema(cls): + return io.Schema( + node_id="LoadVideoTextDataSetFromFolder", + display_name="Load Video and Text Dataset from Folder", + category="dataset", + is_experimental=True, + inputs=[ + io.Combo.Input( + "folder", + options=folder_paths.get_input_subfolders(), + tooltip="The folder containing video files and .txt captions.", + ), + io.Int.Input( + "max_frames", + default=0, + min=0, + max=99999, + tooltip="Maximum frames to load per video (0 = all frames).", + ), + io.Int.Input( + "frame_stride", + default=1, + min=1, + max=1000, + tooltip="Sample every Nth frame (1 = every frame).", + ), + io.Int.Input( + "start_frame", + default=0, + min=0, + max=99999, + tooltip="Frame index to start loading from.", + ), + ], + outputs=[ + io.Image.Output( + display_name="videos", + is_output_list=True, + tooltip="List of video tensors, each [T, H, W, C].", + ), + io.String.Output( + display_name="texts", + is_output_list=True, + tooltip="List of text captions.", + ), + ], + ) + + @classmethod + def execute(cls, folder, max_frames, frame_stride, start_frame): + sub_input_dir = os.path.join(folder_paths.get_input_directory(), folder) + + video_files = [] + for item in sorted(os.listdir(sub_input_dir)): + path = os.path.join(sub_input_dir, item) + if any(item.lower().endswith(ext) for ext in VALID_VIDEO_EXTENSIONS): + video_files.append(path) + elif os.path.isdir(path): + # Support kohya-ss/sd-scripts folder structure: {repeat}_{desc}/ + repeat = 1 + if item.split("_")[0].isdigit(): + repeat = int(item.split("_")[0]) + video_files.extend([ + os.path.join(path, f) + for f in sorted(os.listdir(path)) + if any(f.lower().endswith(ext) for ext in VALID_VIDEO_EXTENSIONS) + ] * repeat) + + if not video_files: + raise ValueError(f"No video files found in {sub_input_dir}") + + # Load captions (same name as video but .txt) + captions = [] + for vf in video_files: + caption_path = os.path.splitext(vf)[0] + ".txt" + if os.path.exists(caption_path): + with open(caption_path, "r", encoding="utf-8") as f: + captions.append(f.read().strip()) + else: + captions.append("") + + # Load videos + output_videos = [] + for vf in video_files: + frames = load_video_frames(vf, max_frames, frame_stride, start_frame) + output_videos.append(frames) + + logging.info(f"Loaded {len(output_videos)} videos with captions from {sub_input_dir}") + return io.NodeOutput(output_videos, captions) + + def save_images_to_folder(image_list, output_dir, prefix="image"): """Utility function to save a list of image tensors to disk. @@ -418,7 +622,15 @@ class ImageProcessingNode(io.ComfyNode): @classmethod def execute(cls, images, **kwargs): - """Execute the node. Routes to _process or _group_process based on mode.""" + """Execute the node. Routes to _process or _group_process based on mode. + + For individual processing (_process), automatically handles multi-frame + inputs (video tensors [T, H, W, C]) by applying _process per-frame and + concatenating the results. This allows all spatial transform nodes to + work with video without modification. Nodes that natively handle batched + tensors (e.g. pure tensor math) can set per_frame_process = False to + skip the per-frame loop. + """ is_group = cls._detect_processing_mode() # Extract scalar values from lists for parameters @@ -434,7 +646,16 @@ class ImageProcessingNode(io.ComfyNode): result = cls._group_process(images, **params) else: # Individual processing: images is single item, call _process - result = cls._process(images, **params) + # Auto-loop over frames for multi-frame inputs (video [T, H, W, C]) + # so that PIL-based spatial transforms work per-frame automatically. + if images.shape[0] > 1 and getattr(cls, 'per_frame_process', True): + results = [] + for i in range(images.shape[0]): + frame_result = cls._process(images[i:i + 1], **params) + results.append(frame_result) + result = torch.cat(results, dim=0) + else: + result = cls._process(images, **params) return io.NodeOutput(result) @@ -736,6 +957,7 @@ class NormalizeImagesNode(ImageProcessingNode): node_id = "NormalizeImages" display_name = "Normalize Images" description = "Normalize images using mean and standard deviation." + per_frame_process = False # Pure tensor math, handles any batch size extra_inputs = [ io.Float.Input( "mean", @@ -764,6 +986,7 @@ class AdjustBrightnessNode(ImageProcessingNode): node_id = "AdjustBrightness" display_name = "Adjust Brightness" description = "Adjust brightness of all images." + per_frame_process = False # Pure tensor math, handles any batch size extra_inputs = [ io.Float.Input( "factor", @@ -783,6 +1006,7 @@ class AdjustContrastNode(ImageProcessingNode): node_id = "AdjustContrast" display_name = "Adjust Contrast" description = "Adjust contrast of all images." + per_frame_process = False # Pure tensor math, handles any batch size extra_inputs = [ io.Float.Input( "factor", @@ -860,6 +1084,243 @@ class ShuffleImageTextDatasetNode(io.ComfyNode): return io.NodeOutput(shuffled_images, shuffled_texts) +# ========== Video Processing Nodes ========== + + +class VideoFrameSampleNode(io.ComfyNode): + """Sample a fixed number of frames from a video using various strategies.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="VideoFrameSample", + display_name="Video Frame Sample", + category="dataset/video", + is_experimental=True, + inputs=[ + io.Image.Input("video", tooltip="Video tensor [T, H, W, C]."), + io.Int.Input( + "num_frames", + default=16, + min=1, + max=9999, + tooltip="Number of frames to sample.", + ), + io.Combo.Input( + "strategy", + options=["uniform", "head", "tail", "random"], + default="uniform", + tooltip="uniform: evenly spaced, head: first N, tail: last N, random: random sorted.", + ), + io.Int.Input( + "seed", + default=0, + min=0, + max=0xFFFFFFFFFFFFFFFF, + tooltip="Random seed (only used with 'random' strategy).", + ), + ], + outputs=[ + io.Image.Output(display_name="video", tooltip="Sampled video [N, H, W, C]."), + ], + ) + + @classmethod + def execute(cls, video, num_frames, strategy, seed): + total_frames = video.shape[0] + num_frames = min(num_frames, total_frames) + + if strategy == "head": + indices = list(range(num_frames)) + elif strategy == "tail": + indices = list(range(total_frames - num_frames, total_frames)) + elif strategy == "uniform": + if num_frames == 1: + indices = [total_frames // 2] + else: + indices = [round(i * (total_frames - 1) / (num_frames - 1)) for i in range(num_frames)] + elif strategy == "random": + rng = np.random.RandomState(seed % (2**32 - 1)) + indices = sorted(rng.choice(total_frames, size=num_frames, replace=False).tolist()) + else: + raise ValueError(f"Unknown strategy: {strategy}") + + return io.NodeOutput(video[indices]) + + +class VideoTemporalCropNode(io.ComfyNode): + """Crop a continuous range of frames from a video.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="VideoTemporalCrop", + display_name="Video Temporal Crop", + category="dataset/video", + is_experimental=True, + inputs=[ + io.Image.Input("video", tooltip="Video tensor [T, H, W, C]."), + io.Int.Input( + "start_frame", + default=0, + min=0, + max=99999, + tooltip="Starting frame index.", + ), + io.Int.Input( + "length", + default=16, + min=1, + max=99999, + tooltip="Number of frames to keep.", + ), + ], + outputs=[ + io.Image.Output(display_name="video", tooltip="Cropped video [length, H, W, C]."), + ], + ) + + @classmethod + def execute(cls, video, start_frame, length): + total_frames = video.shape[0] + start_frame = min(start_frame, max(total_frames - 1, 0)) + end_frame = min(start_frame + length, total_frames) + return io.NodeOutput(video[start_frame:end_frame]) + + +class VideoRandomTemporalCropNode(io.ComfyNode): + """Randomly crop a continuous range of frames from a video (for data augmentation).""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="VideoRandomTemporalCrop", + display_name="Video Random Temporal Crop", + category="dataset/video", + is_experimental=True, + inputs=[ + io.Image.Input("video", tooltip="Video tensor [T, H, W, C]."), + io.Int.Input( + "length", + default=16, + min=1, + max=99999, + tooltip="Number of frames to keep.", + ), + io.Int.Input( + "seed", + default=0, + min=0, + max=0xFFFFFFFFFFFFFFFF, + tooltip="Random seed.", + ), + ], + outputs=[ + io.Image.Output(display_name="video", tooltip="Cropped video [length, H, W, C]."), + ], + ) + + @classmethod + def execute(cls, video, length, seed): + total_frames = video.shape[0] + length = min(length, total_frames) + max_start = total_frames - length + rng = np.random.RandomState(seed % (2**32 - 1)) + start = rng.randint(0, max_start + 1) if max_start > 0 else 0 + return io.NodeOutput(video[start:start + length]) + + +class ShuffleVideoDatasetNode(ImageProcessingNode): + """Randomly shuffle the order of videos in the dataset.""" + + node_id = "ShuffleVideoDataset" + display_name = "Shuffle Video Dataset" + description = "Randomly shuffle the order of videos in the dataset." + is_group_process = True + extra_inputs = [ + io.Int.Input( + "seed", default=0, min=0, max=0xFFFFFFFFFFFFFFFF, tooltip="Random seed." + ), + ] + + @classmethod + def define_schema(cls): + return io.Schema( + node_id=cls.node_id, + display_name=cls.display_name, + category="dataset/video", + is_experimental=True, + is_input_list=True, + inputs=[ + io.Image.Input("images", tooltip="List of videos to shuffle."), + io.Int.Input( + "seed", default=0, min=0, max=0xFFFFFFFFFFFFFFFF, tooltip="Random seed." + ), + ], + outputs=[ + io.Image.Output( + display_name="videos", + is_output_list=True, + tooltip="Shuffled videos", + ), + ], + ) + + @classmethod + def execute(cls, images, seed): + seed = seed[0] if isinstance(seed, list) else seed + np.random.seed(seed % (2**32 - 1)) + indices = np.random.permutation(len(images)) + return io.NodeOutput([images[i] for i in indices]) + + +class ShuffleVideoTextDatasetNode(io.ComfyNode): + """Shuffle videos and their captions together, preserving pairs.""" + + @classmethod + def define_schema(cls): + return io.Schema( + node_id="ShuffleVideoTextDataset", + display_name="Shuffle Video-Text Dataset", + category="dataset/video", + is_experimental=True, + is_input_list=True, + inputs=[ + io.Image.Input("videos", tooltip="List of videos to shuffle."), + io.String.Input("texts", tooltip="List of texts to shuffle."), + io.Int.Input( + "seed", + default=0, + min=0, + max=0xFFFFFFFFFFFFFFFF, + tooltip="Random seed.", + ), + ], + outputs=[ + io.Image.Output( + display_name="videos", + is_output_list=True, + tooltip="Shuffled videos", + ), + io.String.Output( + display_name="texts", + is_output_list=True, + tooltip="Shuffled texts", + ), + ], + ) + + @classmethod + def execute(cls, videos, texts, seed): + seed = seed[0] if isinstance(seed, list) else seed + np.random.seed(seed % (2**32 - 1)) + indices = np.random.permutation(len(videos)) + return io.NodeOutput( + [videos[i] for i in indices], + [texts[i] for i in indices], + ) + + # ========== Text Transform Nodes ========== @@ -1502,7 +1963,10 @@ class DatasetExtension(ComfyExtension): LoadImageTextDataSetFromFolderNode, SaveImageDataSetToFolderNode, SaveImageTextDataSetToFolderNode, - # Image transform nodes + # Video data loading nodes + LoadVideoDataSetFromFolderNode, + LoadVideoTextDataSetFromFolderNode, + # Image transform nodes (auto-handle video via per-frame processing) ResizeImagesByShorterEdgeNode, ResizeImagesByLongerEdgeNode, CenterCropImagesNode, @@ -1512,6 +1976,12 @@ class DatasetExtension(ComfyExtension): AdjustContrastNode, ShuffleDatasetNode, ShuffleImageTextDatasetNode, + # Video processing nodes + VideoFrameSampleNode, + VideoTemporalCropNode, + VideoRandomTemporalCropNode, + ShuffleVideoDatasetNode, + ShuffleVideoTextDatasetNode, # Text transform nodes TextToLowercaseNode, TextToUppercaseNode, From 025bce5ab6137bc511899b621e382c12da132fb2 Mon Sep 17 00:00:00 2001 From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com> Date: Tue, 28 Apr 2026 01:12:57 +0800 Subject: [PATCH 2/4] Ensure train node support real 5D tensor data --- comfy_extras/nodes_train.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/comfy_extras/nodes_train.py b/comfy_extras/nodes_train.py index 0616dfc2d..d6a062d06 100644 --- a/comfy_extras/nodes_train.py +++ b/comfy_extras/nodes_train.py @@ -914,10 +914,11 @@ def _run_training_loop( """ sigmas = torch.tensor(range(num_images)) noise = comfy_extras.nodes_custom_sampler.Noise_RandomNoise(seed) + ndim = latents[0].ndim if bucket_mode: # Use first bucket's first latent as dummy for guider - dummy_latent = latents[0][:1].repeat(num_images, 1, 1, 1) + dummy_latent = latents[0][:1].repeat(num_images, *[1]*(ndim-1)) guider.sample( noise.generate_noise({"samples": dummy_latent}), dummy_latent, @@ -927,7 +928,7 @@ def _run_training_loop( ) elif multi_res: # use first latent as dummy latent if multi_res - latents = latents[0].repeat(num_images, 1, 1, 1) + latents = latents[0].repeat(num_images, *[1]*(ndim-1)) guider.sample( noise.generate_noise({"samples": latents}), latents, From c7d29a42ba2b6ff3161db850c89cf6771f0f90d3 Mon Sep 17 00:00:00 2001 From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com> Date: Tue, 9 Jun 2026 12:19:06 +0800 Subject: [PATCH 3/4] refactor(video dataset): lazy video loading and frame-selective decode - Replace eager load_video_frames() with _decode_selected_frames() that opens the container with `with av.open(...)` (no resource leak) and decodes only the requested frame indices. - Video loader nodes now emit lazy VideoFromFile references; sampling and temporal-crop nodes operate lazily / decode only selected frames. --- comfy_extras/nodes_dataset.py | 227 +++++++++++++--------------------- 1 file changed, 85 insertions(+), 142 deletions(-) diff --git a/comfy_extras/nodes_dataset.py b/comfy_extras/nodes_dataset.py index 089d37dbc..d70a14ccc 100644 --- a/comfy_extras/nodes_dataset.py +++ b/comfy_extras/nodes_dataset.py @@ -10,7 +10,7 @@ from typing_extensions import override import folder_paths import node_helpers -from comfy_api.latest import ComfyExtension, io +from comfy_api.latest import ComfyExtension, io, Input, InputImpl, Types def load_and_process_images(image_files, input_dir): @@ -46,44 +46,33 @@ def load_and_process_images(image_files, input_dir): VALID_VIDEO_EXTENSIONS = [".mp4", ".avi", ".mov", ".webm", ".mkv", ".flv"] -def load_video_frames(video_path, max_frames=0, frame_stride=1, start_frame=0): - """Load video file and return frames as a tensor. +def _decode_selected_frames(video: Input.Video, indices: list[int]) -> Input.Video: + """Decode only the requested frame indices from a video. - Args: - video_path: Path to the video file - max_frames: Maximum number of frames to load (0 = all) - frame_stride: Sample every Nth frame - start_frame: Frame index to start from - - Returns: - torch.Tensor: Video frames as [T, H, W, C] float32 tensor in [0, 1] + Opens the underlying container once, decodes frames in presentation order, + keeps only the ones whose index is in ``indices``, and returns the result + wrapped in a VideoFromComponents so it still satisfies the VideoInput + contract for downstream nodes. """ - container = av.open(video_path) - stream = container.streams.video[0] + indices_sorted = sorted(set(indices)) + max_idx = indices_sorted[-1] + source = video.get_stream_source() - frames = [] - frame_idx = 0 - for frame in container.decode(stream): - if frame_idx < start_frame: - frame_idx += 1 - continue - if (frame_idx - start_frame) % frame_stride != 0: - frame_idx += 1 - continue - if max_frames > 0 and len(frames) >= max_frames: - break + frames_by_idx: dict[int, torch.Tensor] = {} + with av.open(source, mode="r") as container: + stream = container.streams.video[0] + wanted = set(indices_sorted) + for frame_idx, frame in enumerate(container.decode(stream)): + if frame_idx in wanted: + img = frame.to_ndarray(format="rgb24") + frames_by_idx[frame_idx] = torch.from_numpy(img.copy()).float() / 255.0 + if frame_idx >= max_idx: + break - img = frame.to_ndarray(format='rgb24') - img_tensor = torch.from_numpy(img.copy()).float() / 255.0 - frames.append(img_tensor) - frame_idx += 1 - - container.close() - - if not frames: - raise ValueError(f"No frames could be extracted from {video_path}") - - return torch.stack(frames) # [T, H, W, C] + stacked = torch.stack([frames_by_idx[i] for i in indices]) + return InputImpl.VideoFromComponents( + Types.VideoComponents(images=stacked, frame_rate=video.get_frame_rate()) + ) class LoadImageDataSetFromFolderNode(io.ComfyNode): @@ -211,39 +200,18 @@ class LoadVideoDataSetFromFolderNode(io.ComfyNode): options=folder_paths.get_input_subfolders(), tooltip="The folder containing video files.", ), - io.Int.Input( - "max_frames", - default=0, - min=0, - max=99999, - tooltip="Maximum frames to load per video (0 = all frames).", - ), - io.Int.Input( - "frame_stride", - default=1, - min=1, - max=1000, - tooltip="Sample every Nth frame (1 = every frame).", - ), - io.Int.Input( - "start_frame", - default=0, - min=0, - max=99999, - tooltip="Frame index to start loading from.", - ), ], outputs=[ - io.Image.Output( + io.Video.Output( display_name="videos", is_output_list=True, - tooltip="List of video tensors, each [T, H, W, C].", + tooltip="Lazy video references; frames are decoded only when needed downstream.", ), ], ) @classmethod - def execute(cls, folder, max_frames, frame_stride, start_frame): + def execute(cls, folder): sub_input_dir = os.path.join(folder_paths.get_input_directory(), folder) video_files = sorted([ f for f in os.listdir(sub_input_dir) @@ -253,15 +221,9 @@ class LoadVideoDataSetFromFolderNode(io.ComfyNode): if not video_files: raise ValueError(f"No video files found in {sub_input_dir}") - output_videos = [] - for file in video_files: - video_path = os.path.join(sub_input_dir, file) - frames = load_video_frames(video_path, max_frames, frame_stride, start_frame) - output_videos.append(frames) - logging.info(f"Loaded {file}: {frames.shape[0]} frames, {frames.shape[1]}x{frames.shape[2]}") - - logging.info(f"Loaded {len(output_videos)} videos from {sub_input_dir}") - return io.NodeOutput(output_videos) + videos = [InputImpl.VideoFromFile(os.path.join(sub_input_dir, f)) for f in video_files] + logging.info(f"Loaded {len(videos)} lazy video references from {sub_input_dir}") + return io.NodeOutput(videos) class LoadVideoTextDataSetFromFolderNode(io.ComfyNode): @@ -278,33 +240,12 @@ class LoadVideoTextDataSetFromFolderNode(io.ComfyNode): options=folder_paths.get_input_subfolders(), tooltip="The folder containing video files and .txt captions.", ), - io.Int.Input( - "max_frames", - default=0, - min=0, - max=99999, - tooltip="Maximum frames to load per video (0 = all frames).", - ), - io.Int.Input( - "frame_stride", - default=1, - min=1, - max=1000, - tooltip="Sample every Nth frame (1 = every frame).", - ), - io.Int.Input( - "start_frame", - default=0, - min=0, - max=99999, - tooltip="Frame index to start loading from.", - ), ], outputs=[ - io.Image.Output( + io.Video.Output( display_name="videos", is_output_list=True, - tooltip="List of video tensors, each [T, H, W, C].", + tooltip="Lazy video references; frames are decoded only when needed downstream.", ), io.String.Output( display_name="texts", @@ -315,7 +256,7 @@ class LoadVideoTextDataSetFromFolderNode(io.ComfyNode): ) @classmethod - def execute(cls, folder, max_frames, frame_stride, start_frame): + def execute(cls, folder): sub_input_dir = os.path.join(folder_paths.get_input_directory(), folder) video_files = [] @@ -337,7 +278,6 @@ class LoadVideoTextDataSetFromFolderNode(io.ComfyNode): if not video_files: raise ValueError(f"No video files found in {sub_input_dir}") - # Load captions (same name as video but .txt) captions = [] for vf in video_files: caption_path = os.path.splitext(vf)[0] + ".txt" @@ -347,14 +287,9 @@ class LoadVideoTextDataSetFromFolderNode(io.ComfyNode): else: captions.append("") - # Load videos - output_videos = [] - for vf in video_files: - frames = load_video_frames(vf, max_frames, frame_stride, start_frame) - output_videos.append(frames) - - logging.info(f"Loaded {len(output_videos)} videos with captions from {sub_input_dir}") - return io.NodeOutput(output_videos, captions) + videos = [InputImpl.VideoFromFile(vf) for vf in video_files] + logging.info(f"Loaded {len(videos)} lazy video references with captions from {sub_input_dir}") + return io.NodeOutput(videos, captions) def save_images_to_folder(image_list, output_dir, prefix="image"): @@ -1088,7 +1023,12 @@ class ShuffleImageTextDatasetNode(io.ComfyNode): class VideoFrameSampleNode(io.ComfyNode): - """Sample a fixed number of frames from a video using various strategies.""" + """Sample a fixed number of frames from a video using various strategies. + + For contiguous strategies ("head"/"tail") the result is a fully lazy + VideoInput (no frames decoded). For non-contiguous strategies + ("uniform"/"random") only the selected indices are decoded. + """ @classmethod def define_schema(cls): @@ -1098,7 +1038,7 @@ class VideoFrameSampleNode(io.ComfyNode): category="dataset/video", is_experimental=True, inputs=[ - io.Image.Input("video", tooltip="Video tensor [T, H, W, C]."), + io.Video.Input("video", tooltip="Input video."), io.Int.Input( "num_frames", default=16, @@ -1121,20 +1061,27 @@ class VideoFrameSampleNode(io.ComfyNode): ), ], outputs=[ - io.Image.Output(display_name="video", tooltip="Sampled video [N, H, W, C]."), + io.Video.Output(display_name="video", tooltip="Sampled video."), ], ) @classmethod def execute(cls, video, num_frames, strategy, seed): - total_frames = video.shape[0] + total_frames = video.get_frame_count() num_frames = min(num_frames, total_frames) + fps = float(video.get_frame_rate()) if strategy == "head": - indices = list(range(num_frames)) - elif strategy == "tail": - indices = list(range(total_frames - num_frames, total_frames)) - elif strategy == "uniform": + return io.NodeOutput( + video.as_trimmed(0.0, num_frames / fps, strict_duration=False) + ) + if strategy == "tail": + start_t = (total_frames - num_frames) / fps + return io.NodeOutput( + video.as_trimmed(start_t, num_frames / fps, strict_duration=False) + ) + + if strategy == "uniform": if num_frames == 1: indices = [total_frames // 2] else: @@ -1145,11 +1092,11 @@ class VideoFrameSampleNode(io.ComfyNode): else: raise ValueError(f"Unknown strategy: {strategy}") - return io.NodeOutput(video[indices]) + return io.NodeOutput(_decode_selected_frames(video, indices)) class VideoTemporalCropNode(io.ComfyNode): - """Crop a continuous range of frames from a video.""" + """Crop a continuous range of frames from a video (fully lazy).""" @classmethod def define_schema(cls): @@ -1159,7 +1106,7 @@ class VideoTemporalCropNode(io.ComfyNode): category="dataset/video", is_experimental=True, inputs=[ - io.Image.Input("video", tooltip="Video tensor [T, H, W, C]."), + io.Video.Input("video", tooltip="Input video."), io.Int.Input( "start_frame", default=0, @@ -1176,20 +1123,23 @@ class VideoTemporalCropNode(io.ComfyNode): ), ], outputs=[ - io.Image.Output(display_name="video", tooltip="Cropped video [length, H, W, C]."), + io.Video.Output(display_name="video", tooltip="Cropped video (lazy)."), ], ) @classmethod def execute(cls, video, start_frame, length): - total_frames = video.shape[0] + total_frames = video.get_frame_count() + fps = float(video.get_frame_rate()) start_frame = min(start_frame, max(total_frames - 1, 0)) - end_frame = min(start_frame + length, total_frames) - return io.NodeOutput(video[start_frame:end_frame]) + length = min(length, total_frames - start_frame) + return io.NodeOutput( + video.as_trimmed(start_frame / fps, length / fps, strict_duration=False) + ) class VideoRandomTemporalCropNode(io.ComfyNode): - """Randomly crop a continuous range of frames from a video (for data augmentation).""" + """Randomly crop a continuous range of frames from a video (fully lazy).""" @classmethod def define_schema(cls): @@ -1199,7 +1149,7 @@ class VideoRandomTemporalCropNode(io.ComfyNode): category="dataset/video", is_experimental=True, inputs=[ - io.Image.Input("video", tooltip="Video tensor [T, H, W, C]."), + io.Video.Input("video", tooltip="Input video."), io.Int.Input( "length", default=16, @@ -1216,49 +1166,42 @@ class VideoRandomTemporalCropNode(io.ComfyNode): ), ], outputs=[ - io.Image.Output(display_name="video", tooltip="Cropped video [length, H, W, C]."), + io.Video.Output(display_name="video", tooltip="Cropped video (lazy)."), ], ) @classmethod def execute(cls, video, length, seed): - total_frames = video.shape[0] + total_frames = video.get_frame_count() + fps = float(video.get_frame_rate()) length = min(length, total_frames) max_start = total_frames - length rng = np.random.RandomState(seed % (2**32 - 1)) start = rng.randint(0, max_start + 1) if max_start > 0 else 0 - return io.NodeOutput(video[start:start + length]) + return io.NodeOutput( + video.as_trimmed(start / fps, length / fps, strict_duration=False) + ) -class ShuffleVideoDatasetNode(ImageProcessingNode): +class ShuffleVideoDatasetNode(io.ComfyNode): """Randomly shuffle the order of videos in the dataset.""" - node_id = "ShuffleVideoDataset" - display_name = "Shuffle Video Dataset" - description = "Randomly shuffle the order of videos in the dataset." - is_group_process = True - extra_inputs = [ - io.Int.Input( - "seed", default=0, min=0, max=0xFFFFFFFFFFFFFFFF, tooltip="Random seed." - ), - ] - @classmethod def define_schema(cls): return io.Schema( - node_id=cls.node_id, - display_name=cls.display_name, + node_id="ShuffleVideoDataset", + display_name="Shuffle Video Dataset", category="dataset/video", is_experimental=True, is_input_list=True, inputs=[ - io.Image.Input("images", tooltip="List of videos to shuffle."), + io.Video.Input("videos", tooltip="List of videos to shuffle."), io.Int.Input( "seed", default=0, min=0, max=0xFFFFFFFFFFFFFFFF, tooltip="Random seed." ), ], outputs=[ - io.Image.Output( + io.Video.Output( display_name="videos", is_output_list=True, tooltip="Shuffled videos", @@ -1267,11 +1210,11 @@ class ShuffleVideoDatasetNode(ImageProcessingNode): ) @classmethod - def execute(cls, images, seed): + def execute(cls, videos, seed): seed = seed[0] if isinstance(seed, list) else seed np.random.seed(seed % (2**32 - 1)) - indices = np.random.permutation(len(images)) - return io.NodeOutput([images[i] for i in indices]) + indices = np.random.permutation(len(videos)) + return io.NodeOutput([videos[i] for i in indices]) class ShuffleVideoTextDatasetNode(io.ComfyNode): @@ -1286,7 +1229,7 @@ class ShuffleVideoTextDatasetNode(io.ComfyNode): is_experimental=True, is_input_list=True, inputs=[ - io.Image.Input("videos", tooltip="List of videos to shuffle."), + io.Video.Input("videos", tooltip="List of videos to shuffle."), io.String.Input("texts", tooltip="List of texts to shuffle."), io.Int.Input( "seed", @@ -1297,7 +1240,7 @@ class ShuffleVideoTextDatasetNode(io.ComfyNode): ), ], outputs=[ - io.Image.Output( + io.Video.Output( display_name="videos", is_output_list=True, tooltip="Shuffled videos", @@ -1976,7 +1919,7 @@ class DatasetExtension(ComfyExtension): AdjustContrastNode, ShuffleDatasetNode, ShuffleImageTextDatasetNode, - # Video processing nodes + # Video processing nodes (lazy VideoInput in/out) VideoFrameSampleNode, VideoTemporalCropNode, VideoRandomTemporalCropNode, From 133a872cc8ca033e87e6e22b4b6d16c9a0426b32 Mon Sep 17 00:00:00 2001 From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com> Date: Tue, 9 Jun 2026 12:26:45 +0800 Subject: [PATCH 4/4] style(video nodes): align naming/category with image nodes (PR #13588 review) Apply the same naming convention the image nodes adopted, per @alexisrolland's review on PR #13588: - Load Video (from Folder), Load Video-Text (from Folder): category "video", add search_aliases + description. - Sample Video Frame (was "Video Frame Sample"): category "video". - Crop Video (Temporal) / Crop Video (Temporal Random): category "video/transform", verb-first names, search_aliases + description. - Shuffle Videos List (was "Shuffle Video Dataset"): category "video/batch". - Shuffle Pairs of Video-Text: category "dataset/video". All video schemas now carry search_aliases and description to match the image-node conventions. Resolved review threads on PR #13588. Two comments were outdated and intentionally skipped: - CodeRabbit container-leak on load_video_frames(): the function was replaced by _decode_selected_frames() using `with av.open(...)`. - "category=cls.category" on ShuffleVideoDataset: the node is no longer an ImageProcessingNode subclass, so category is set directly. --- comfy_extras/nodes_dataset.py | 40 +++++++++++++++++++++++------------ 1 file changed, 27 insertions(+), 13 deletions(-) diff --git a/comfy_extras/nodes_dataset.py b/comfy_extras/nodes_dataset.py index 3e644b92f..ec3878165 100644 --- a/comfy_extras/nodes_dataset.py +++ b/comfy_extras/nodes_dataset.py @@ -195,8 +195,10 @@ class LoadVideoDataSetFromFolderNode(io.ComfyNode): def define_schema(cls): return io.Schema( node_id="LoadVideoDataSetFromFolder", - display_name="Load Video Dataset from Folder", - category="dataset", + search_aliases=["load folder", "load from folder", "load dataset", "load videos", "import dataset"], + display_name="Load Video (from Folder)", + category="video", + description="Load a dataset of videos from a specified folder and return a list of videos. Supported formats: MP4, AVI, MOV, WEBM, MKV, FLV.", is_experimental=True, inputs=[ io.Combo.Input( @@ -235,8 +237,10 @@ class LoadVideoTextDataSetFromFolderNode(io.ComfyNode): def define_schema(cls): return io.Schema( node_id="LoadVideoTextDataSetFromFolder", - display_name="Load Video and Text Dataset from Folder", - category="dataset", + search_aliases=["load folder", "load from folder", "load dataset", "load videos", "import dataset"], + display_name="Load Video-Text (from Folder)", + category="video", + description="Load a dataset of pairs of videos and text captions from a specified folder and return them as a list. Supported formats: MP4, AVI, MOV, WEBM, MKV, FLV.", is_experimental=True, inputs=[ io.Combo.Input( @@ -1109,8 +1113,10 @@ class VideoFrameSampleNode(io.ComfyNode): def define_schema(cls): return io.Schema( node_id="VideoFrameSample", - display_name="Video Frame Sample", - category="dataset/video", + search_aliases=["sample frames", "extract frames"], + display_name="Sample Video Frame", + category="video", + description="Sample a fixed number of frames from a video using various strategies.", is_experimental=True, inputs=[ io.Video.Input("video", tooltip="Input video."), @@ -1177,8 +1183,10 @@ class VideoTemporalCropNode(io.ComfyNode): def define_schema(cls): return io.Schema( node_id="VideoTemporalCrop", - display_name="Video Temporal Crop", - category="dataset/video", + search_aliases=["crop", "crop video", "temporal crop", "truncate video"], + display_name="Crop Video (Temporal)", + category="video/transform", + description="Crop a continuous range of frames from a video.", is_experimental=True, inputs=[ io.Video.Input("video", tooltip="Input video."), @@ -1220,8 +1228,10 @@ class VideoRandomTemporalCropNode(io.ComfyNode): def define_schema(cls): return io.Schema( node_id="VideoRandomTemporalCrop", - display_name="Video Random Temporal Crop", - category="dataset/video", + search_aliases=["crop", "crop video", "temporal crop", "truncate video", "random crop"], + display_name="Crop Video (Temporal Random)", + category="video/transform", + description="Randomly crop a continuous range of frames from a video.", is_experimental=True, inputs=[ io.Video.Input("video", tooltip="Input video."), @@ -1265,8 +1275,10 @@ class ShuffleVideoDatasetNode(io.ComfyNode): def define_schema(cls): return io.Schema( node_id="ShuffleVideoDataset", - display_name="Shuffle Video Dataset", - category="dataset/video", + search_aliases=["shuffle", "randomize", "mix"], + display_name="Shuffle Videos List", + category="video/batch", + description="Randomly shuffle the order of videos in a list.", is_experimental=True, is_input_list=True, inputs=[ @@ -1299,8 +1311,10 @@ class ShuffleVideoTextDatasetNode(io.ComfyNode): def define_schema(cls): return io.Schema( node_id="ShuffleVideoTextDataset", - display_name="Shuffle Video-Text Dataset", + search_aliases=["shuffle", "randomize", "mix"], + display_name="Shuffle Pairs of Video-Text", category="dataset/video", + description="Randomly shuffle the order of pairs of video-text in a list.", is_experimental=True, is_input_list=True, inputs=[