|
| 1 | +import os |
| 2 | +import re |
| 3 | +import cv2 |
| 4 | +import torch |
| 5 | + |
| 6 | + |
| 7 | +class LoadVideosFromFolderSimple: |
| 8 | + VIDEO_EXTENSIONS = ['webm', 'mp4', 'mkv', 'gif', 'mov'] |
| 9 | + |
| 10 | + @classmethod |
| 11 | + def INPUT_TYPES(cls): |
| 12 | + return { |
| 13 | + "required": { |
| 14 | + "folder_path": ("STRING", {"default": ""}), |
| 15 | + "debug": ("BOOLEAN", { |
| 16 | + "default": False, |
| 17 | + "tooltip": "Log some details to the console" |
| 18 | + }), |
| 19 | + }, |
| 20 | + "optional": { |
| 21 | + "meta_batch": ("VHS_BatchManager",), |
| 22 | + }, |
| 23 | + "hidden": { |
| 24 | + "unique_id": "UNIQUE_ID", |
| 25 | + }, |
| 26 | + } |
| 27 | + |
| 28 | + RETURN_TYPES = ["IMAGE"] |
| 29 | + RETURN_NAMES = ["images"] |
| 30 | + FUNCTION = "load_videos" |
| 31 | + CATEGORY = "video/utility" |
| 32 | + DESCRIPTION = """ |
| 33 | + Load all videos from a folder, concatenated into a single image batch. |
| 34 | + Optionally connect a VideoHelperSuite Meta Batch Manager to process large collections in smaller RAM-safe chunks. See VHS Meta Batch Manager documentation for more information. |
| 35 | + - Formats: webm, mp4, mkv, gif, mov |
| 36 | + - All videos must have identical resolution |
| 37 | + """ |
| 38 | + |
| 39 | + def load_videos(self, folder_path, debug, meta_batch=None, unique_id=None): |
| 40 | + folder_path = folder_path.strip().strip('"').strip("'") |
| 41 | + |
| 42 | + if not os.path.isdir(folder_path): |
| 43 | + raise ValueError(f"Folder does not exist: {folder_path}") |
| 44 | + |
| 45 | + video_files = self._get_video_files(folder_path) |
| 46 | + |
| 47 | + if not video_files: |
| 48 | + raise ValueError( |
| 49 | + f"No video files found in {folder_path}\n" |
| 50 | + f"Supported formats: {', '.join(self.VIDEO_EXTENSIONS)}" |
| 51 | + ) |
| 52 | + |
| 53 | + if meta_batch is None: |
| 54 | + return self._load_all(video_files, folder_path, debug) |
| 55 | + else: |
| 56 | + return self._load_batched(video_files, folder_path, debug, meta_batch, unique_id) |
| 57 | + |
| 58 | + def _load_all(self, video_files, folder_path, debug): |
| 59 | + """Load all videos at once into a single tensor.""" |
| 60 | + if debug: |
| 61 | + print(f"Loading {len(video_files)} videos from {folder_path}") |
| 62 | + |
| 63 | + all_frames = [] |
| 64 | + expected_shape = None |
| 65 | + |
| 66 | + for idx, video_path in enumerate(video_files): |
| 67 | + if debug: |
| 68 | + print(f"[{idx+1}/{len(video_files)}]: {os.path.basename(video_path)}", end=" ... ") |
| 69 | + |
| 70 | + frames = self._load_video_frames(video_path) |
| 71 | + expected_shape = self._check_resolution(frames, expected_shape, video_path) |
| 72 | + all_frames.append(frames) |
| 73 | + |
| 74 | + if debug: |
| 75 | + print(f"{frames.shape[0]} frames") |
| 76 | + |
| 77 | + if debug: |
| 78 | + print(f"\nConcatenating {len(video_files)} videos...") |
| 79 | + output = torch.cat(all_frames, dim=0) |
| 80 | + |
| 81 | + if debug: |
| 82 | + print(f"Done\n") |
| 83 | + return (output,) |
| 84 | + |
| 85 | + def _load_batched(self, video_files, folder_path, debug, meta_batch, unique_id): |
| 86 | + """ |
| 87 | + Load frames in chunks coordinated by the VHS BatchManager. |
| 88 | + The BatchManager drives re-execution of the workflow. Each pass |
| 89 | + picks up where the generator left off, yields up to frames_per_batch |
| 90 | + frames, then returns. The BatchManager re-queues until exhausted. |
| 91 | + """ |
| 92 | + if unique_id not in meta_batch.inputs: |
| 93 | + total_frames = self._count_total_frames(video_files) |
| 94 | + meta_batch.total_frames = min(meta_batch.total_frames, total_frames) |
| 95 | + if debug: |
| 96 | + print(f"[Batched] Starting new generator for {len(video_files)} videos ({total_frames} frames) in {folder_path}") |
| 97 | + meta_batch.inputs[unique_id] = self._frame_generator(video_files, debug) |
| 98 | + |
| 99 | + generator = meta_batch.inputs[unique_id] |
| 100 | + frames_per_batch = meta_batch.frames_per_batch |
| 101 | + |
| 102 | + batch_frames = [] |
| 103 | + expected_shape = None |
| 104 | + frames_collected = 0 |
| 105 | + |
| 106 | + while frames_collected < frames_per_batch: |
| 107 | + try: |
| 108 | + frame_tensor, video_path = next(generator) |
| 109 | + except StopIteration: |
| 110 | + if debug: |
| 111 | + print(f"[Batched] Generator exhausted, cleaning up") |
| 112 | + meta_batch.inputs.pop(unique_id) |
| 113 | + meta_batch.has_closed_inputs = True |
| 114 | + break |
| 115 | + |
| 116 | + expected_shape = self._check_resolution( |
| 117 | + frame_tensor.unsqueeze(0), expected_shape, video_path |
| 118 | + ) |
| 119 | + batch_frames.append(frame_tensor) |
| 120 | + frames_collected += 1 |
| 121 | + |
| 122 | + if not batch_frames: |
| 123 | + raise RuntimeError("Batched loader produced no frames") |
| 124 | + |
| 125 | + output = torch.stack(batch_frames, dim=0) |
| 126 | + |
| 127 | + if debug: |
| 128 | + print(f"[Batched] Yielding {output.shape[0]} frames shape={tuple(output.shape)}") |
| 129 | + |
| 130 | + return (output,) |
| 131 | + |
| 132 | + def _get_video_files(self, folder_path): |
| 133 | + """Return naturally-sorted list of video file paths in folder_path.""" |
| 134 | + video_files = [] |
| 135 | + for f in os.listdir(folder_path): |
| 136 | + full_path = os.path.join(folder_path, f) |
| 137 | + if os.path.isfile(full_path) and self._is_video_file(f): |
| 138 | + video_files.append(full_path) |
| 139 | + |
| 140 | + def natural_sort_key(path): |
| 141 | + return [int(text) if text.isdigit() else text.lower() |
| 142 | + for text in re.split('([0-9]+)', os.path.basename(path))] |
| 143 | + |
| 144 | + return sorted(video_files, key=natural_sort_key) |
| 145 | + |
| 146 | + def _is_video_file(self, filename): |
| 147 | + """Return True if filename has a supported video extension.""" |
| 148 | + ext = filename.split('.')[-1].lower() if '.' in filename else '' |
| 149 | + return ext in self.VIDEO_EXTENSIONS |
| 150 | + |
| 151 | + def _load_video_frames(self, video_path): |
| 152 | + """Load all frames from a video file. Returns [N, H, W, C] float32 tensor.""" |
| 153 | + cap = cv2.VideoCapture(video_path) |
| 154 | + if not cap.isOpened(): |
| 155 | + raise RuntimeError(f"Failed to open video: {video_path}") |
| 156 | + |
| 157 | + frames = [] |
| 158 | + while True: |
| 159 | + ret, frame = cap.read() |
| 160 | + if not ret: |
| 161 | + break |
| 162 | + frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) |
| 163 | + frame_tensor = torch.from_numpy(frame_rgb).float() / 255.0 |
| 164 | + frames.append(frame_tensor) |
| 165 | + |
| 166 | + cap.release() |
| 167 | + |
| 168 | + if not frames: |
| 169 | + raise RuntimeError(f"No frames extracted from {video_path}") |
| 170 | + |
| 171 | + return torch.stack(frames, dim=0) |
| 172 | + |
| 173 | + def _check_resolution(self, frames, expected_shape, video_path): |
| 174 | + """ |
| 175 | + Verify frames match expected_shape (W, H). |
| 176 | + Returns expected_shape, initialising it from frames if not yet set. |
| 177 | + Raises RuntimeError with filename on mismatch. |
| 178 | + """ |
| 179 | + if expected_shape is None: |
| 180 | + return frames.shape[1:3] |
| 181 | + if frames.shape[1:3] != expected_shape: |
| 182 | + raise RuntimeError( |
| 183 | + f"\nResolution mismatch\n" |
| 184 | + f" Expected: {expected_shape[1]}x{expected_shape[0]} (from first video)\n" |
| 185 | + f" Got: {frames.shape[2]}x{frames.shape[1]} in {os.path.basename(video_path)}" |
| 186 | + ) |
| 187 | + return expected_shape |
| 188 | + |
| 189 | + def _count_total_frames(self, video_files): |
| 190 | + """Fast frame count estimate using container metadata. No decoding.""" |
| 191 | + total = 0 |
| 192 | + for video_path in video_files: |
| 193 | + cap = cv2.VideoCapture(video_path) |
| 194 | + if cap.isOpened(): |
| 195 | + total += int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) |
| 196 | + cap.release() |
| 197 | + return total |
| 198 | + |
| 199 | + def _frame_generator(self, video_files, debug): |
| 200 | + """ |
| 201 | + Generator that yields (frame_tensor, video_path) one frame at a time |
| 202 | + across all videos. Keeps at most one video open at a time. |
| 203 | + """ |
| 204 | + for idx, video_path in enumerate(video_files): |
| 205 | + if debug: |
| 206 | + print(f"[Batched] Opening [{idx+1}/{len(video_files)}]: {os.path.basename(video_path)}") |
| 207 | + |
| 208 | + cap = cv2.VideoCapture(video_path) |
| 209 | + if not cap.isOpened(): |
| 210 | + raise RuntimeError(f"Failed to open video: {video_path}") |
| 211 | + |
| 212 | + frame_count = 0 |
| 213 | + while True: |
| 214 | + ret, frame = cap.read() |
| 215 | + if not ret: |
| 216 | + break |
| 217 | + frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) |
| 218 | + frame_tensor = torch.from_numpy(frame_rgb).float() / 255.0 |
| 219 | + frame_count += 1 |
| 220 | + yield frame_tensor, video_path |
| 221 | + |
| 222 | + cap.release() |
| 223 | + |
| 224 | + if debug: |
| 225 | + print(f"[Batched] Finished {os.path.basename(video_path)}: {frame_count} frames") |
| 226 | + |
| 227 | + |
| 228 | +# ComfyUI node registration |
| 229 | +NODE_CLASS_MAPPINGS = { |
| 230 | + "LoadVideosFromFolderSimple": LoadVideosFromFolderSimple, |
| 231 | +} |
| 232 | + |
| 233 | +NODE_DISPLAY_NAME_MAPPINGS = { |
| 234 | + "LoadVideosFromFolderSimple": "🪐 Load Videos From Folder (Simple)", |
| 235 | +} |
0 commit comments