From cc8cfff649b5a358d18ca989f767aaaeb045b40c Mon Sep 17 00:00:00 2001 From: stuttlepress <178523610+stuttlepress@users.noreply.github.com> Date: Fri, 27 Feb 2026 16:16:53 -0600 Subject: [PATCH] Integrated baseline. Still rough around the edges. --- frame_selector.py | 120 +++-------------------- load_videos.py | 76 ++++++++------- wan_vace_batch_context.py | 6 +- wan_vace_extend.py | 10 +- wan_vace_prep.py | 17 ++-- wan_vace_prep_inline.py | 27 +++--- web/visual_frame_selector.js | 178 ++++++++++++++++++++++++----------- 7 files changed, 204 insertions(+), 230 deletions(-) diff --git a/frame_selector.py b/frame_selector.py index e5c3bc5..b6e5f2f 100644 --- a/frame_selector.py +++ b/frame_selector.py @@ -1,37 +1,17 @@ +import os import torch import numpy as np -import os import av import folder_paths -# Video extensions supported for the file selector -VIDEO_EXTENSIONS = ['mp4', 'webm', 'mkv', 'mov', 'avi', 'gif'] - class VisualFrameSelector: + DESCRIPTION = """ + Interactive video frame selection. + - Drag markers to set start/end frames + - Click scrubber to seek + - Use transport buttons for playback """ - Interactive video frame selection with visual preview. - - Load a video and use the visual scrubber to select a range of frames. - The selected frames are output as an IMAGE tensor. - - Frame range is inclusive at both ends, 0-based indexing. - start_frame=0, end_frame=4 outputs 5 frames (0, 1, 2, 3, 4). - - The UI enforces end_frame >= start_frame + 1 (minimum 2-frame selection). - end_frame=0 means "last frame of the video". - """ - - DESCRIPTION = """# Visual Frame Selector - -Interactive video frame selection with visual preview. - -**Controls:** -- Drag markers to set start/end frames -- Click scrubber to seek -- Use transport buttons for playback -- Right-click video for volume control -""" @classmethod def INPUT_TYPES(cls): @@ -47,10 +27,10 @@ Interactive video frame selection with visual preview. } } - RETURN_TYPES = ("IMAGE", "INT", "FLOAT", "INT", "INT", "INT", "AUDIO") - RETURN_NAMES = ("images", "selected_frames", "fps", "total_frames", "start_frame", "end_frame", "audio") + RETURN_TYPES = ("IMAGE", "INT", "INT", "INT", "INT") + RETURN_NAMES = ("images", "selected_frames", "total_frames", "start_frame", "end_frame") FUNCTION = "load_frames" - CATEGORY = "image/video" + CATEGORY = "video/utility" @classmethod def IS_CHANGED(cls, video, start_frame, end_frame, **kwargs): @@ -128,20 +108,16 @@ Interactive video frame selection with visual preview. frames_tensor = torch.stack(frames) - audio = self._extract_audio(video_path, actual_start, actual_end, fps) - return ( frames_tensor, len(frames), - fps, total_frames, actual_start, actual_end, - audio, ) except (ValueError, FileNotFoundError): - raise # Let these propagate cleanly to ComfyUI + raise except Exception as e: raise ValueError(f"Error processing video: {str(e)}") @@ -150,81 +126,11 @@ Interactive video frame selection with visual preview. if container is not None: container.close() - def _extract_audio(self, video_path, start_frame, end_frame, fps): - """Extract audio for the selected frame range using PyAV. - - Returns audio dict compatible with ComfyUI AUDIO type, or None if - no audio track exists. - """ - if fps <= 0: - return None - - try: - container = av.open(video_path) - - if len(container.streams.audio) == 0: - container.close() - return None - - audio_stream = container.streams.audio[0] - sample_rate = audio_stream.rate - channels = audio_stream.channels - - start_time = start_frame / fps - duration = (end_frame - start_frame + 1) / fps - end_time = start_time + duration - - # Seek to start time - if start_time > 0: - target_pts = int(start_time / audio_stream.time_base) - container.seek(target_pts, stream=audio_stream) - - # Decode audio samples - audio_frames = [] - for frame in container.decode(audio=0): - frame_time = float(frame.pts * audio_stream.time_base) if frame.pts is not None else 0.0 - if frame_time > end_time: - break - # Convert to float32 numpy array - arr = frame.to_ndarray() - audio_frames.append(arr) - - container.close() - - if len(audio_frames) == 0: - return None - - # Concatenate all decoded audio - audio_np = np.concatenate(audio_frames, axis=-1) - - # Ensure shape is [channels, samples] - if audio_np.ndim == 1: - audio_np = audio_np.reshape(1, -1) - - # Trim to exact time range - start_sample = int(start_time * sample_rate) - end_sample = int(end_time * sample_rate) - total_samples = audio_np.shape[-1] - - # After seeking, the first decoded frame may start before our target - # Just use what we have, trim to duration - max_samples = end_sample - start_sample - if audio_np.shape[-1] > max_samples: - audio_np = audio_np[:, :max_samples] - - # ComfyUI AUDIO format: {"waveform": tensor [batch, channels, samples], "sample_rate": int} - audio_tensor = torch.from_numpy(audio_np.astype(np.float32)) - audio_tensor = audio_tensor.unsqueeze(0) # Add batch dim - - return {"waveform": audio_tensor, "sample_rate": sample_rate} - - except Exception: - return None NODE_CLASS_MAPPINGS = { - "VisualFrameSelector": VisualFrameSelector, + "VisualFrameSelector": VisualFrameSelector } NODE_DISPLAY_NAME_MAPPINGS = { - "VisualFrameSelector": "🪐 Visual Frame Selector", -} \ No newline at end of file + "VisualFrameSelector": "🪐 Visual Frame Selector" +} diff --git a/load_videos.py b/load_videos.py index 5bf0f58..72a81c5 100644 --- a/load_videos.py +++ b/load_videos.py @@ -1,6 +1,7 @@ import os import re -import cv2 +import av +import numpy as np import torch @@ -25,15 +26,16 @@ class LoadVideosFromFolderSimple: }, } - RETURN_TYPES = ["IMAGE"] - RETURN_NAMES = ["images"] + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("images",) FUNCTION = "load_videos" CATEGORY = "video/utility" DESCRIPTION = """ - Load all videos from a folder, concatenated into a single image batch. - 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. - - Formats: webm, mp4, mkv, gif, mov - - All videos must have identical resolution + Load all videos from a folder, concatenated into + a single image batch. 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. """ def load_videos(self, folder_path, debug, meta_batch=None, unique_id=None): @@ -150,20 +152,19 @@ class LoadVideosFromFolderSimple: def _load_video_frames(self, video_path): """Load all frames from a video file. Returns [N, H, W, C] float32 tensor.""" - cap = cv2.VideoCapture(video_path) - if not cap.isOpened(): - raise RuntimeError(f"Failed to open video: {video_path}") + container = av.open(video_path) + + if len(container.streams.video) == 0: + container.close() + raise RuntimeError(f"No video stream found in: {video_path}") frames = [] - while True: - ret, frame = cap.read() - if not ret: - break - frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - frame_tensor = torch.from_numpy(frame_rgb).float() / 255.0 + for frame in container.decode(video=0): + rgb = frame.to_ndarray(format="rgb24") + frame_tensor = torch.from_numpy(rgb.astype(np.float32) / 255.0) frames.append(frame_tensor) - cap.release() + container.close() if not frames: raise RuntimeError(f"No frames extracted from {video_path}") @@ -190,10 +191,18 @@ class LoadVideosFromFolderSimple: """Fast frame count estimate using container metadata. No decoding.""" total = 0 for video_path in video_files: - cap = cv2.VideoCapture(video_path) - if cap.isOpened(): - total += int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) - cap.release() + try: + container = av.open(video_path) + if len(container.streams.video) > 0: + count = container.streams.video[0].frames + if count > 0: + total += count + else: + # Container doesn't report frame count; decode to count + total += sum(1 for _ in container.decode(video=0)) + container.close() + except Exception: + pass return total def _frame_generator(self, video_files, debug): @@ -205,31 +214,28 @@ class LoadVideosFromFolderSimple: if debug: print(f"[Batched] Opening [{idx+1}/{len(video_files)}]: {os.path.basename(video_path)}") - cap = cv2.VideoCapture(video_path) - if not cap.isOpened(): - raise RuntimeError(f"Failed to open video: {video_path}") + container = av.open(video_path) + + if len(container.streams.video) == 0: + container.close() + raise RuntimeError(f"No video stream found in: {video_path}") frame_count = 0 - while True: - ret, frame = cap.read() - if not ret: - break - frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - frame_tensor = torch.from_numpy(frame_rgb).float() / 255.0 + for frame in container.decode(video=0): + rgb = frame.to_ndarray(format="rgb24") + frame_tensor = torch.from_numpy(rgb.astype(np.float32) / 255.0) frame_count += 1 yield frame_tensor, video_path - cap.release() + container.close() if debug: print(f"[Batched] Finished {os.path.basename(video_path)}: {frame_count} frames") - -# ComfyUI node registration NODE_CLASS_MAPPINGS = { - "LoadVideosFromFolderSimple": LoadVideosFromFolderSimple, + "LoadVideosFromFolderSimple": LoadVideosFromFolderSimple } NODE_DISPLAY_NAME_MAPPINGS = { - "LoadVideosFromFolderSimple": "🪐 Load Videos From Folder (Simple)", + "LoadVideosFromFolderSimple": "🪐 Load Videos From Folder (Simple)" } diff --git a/wan_vace_batch_context.py b/wan_vace_batch_context.py index 5419c36..b29f15a 100644 --- a/wan_vace_batch_context.py +++ b/wan_vace_batch_context.py @@ -30,7 +30,9 @@ class WanVACEBatchContext: RETURN_NAMES = ("work_dir", "workfile_prefix", "video_1_filename", "video_2_filename", "is_first", "is_last") FUNCTION = "setup_context" CATEGORY = "video/VACE" - DESCRIPTION = "Establishes iteration context for batch video processing." + DESCRIPTION = """ + Establishes iteration context for batch video join processing. + """ INPUT_IS_LIST = True def setup_context(self, **kwargs): @@ -89,5 +91,5 @@ NODE_CLASS_MAPPINGS = { } NODE_DISPLAY_NAME_MAPPINGS = { - "WanVACEBatchContext": "🪐 Wan VACE Batch Context" + "WanVACEBatchContext": "🪐 VACE Batch Context" } \ No newline at end of file diff --git a/wan_vace_extend.py b/wan_vace_extend.py index ad69658..acd3449 100644 --- a/wan_vace_extend.py +++ b/wan_vace_extend.py @@ -8,8 +8,8 @@ class WanVACEExtend: "video": ("IMAGE",), "extend_from_idx": ("INT", { "default": -1, - "min": -999999, - "max": 999999, + "min": -1000000, + "max": 1000000, "step": 1, "tooltip": "Frame index to extend from. Negative values count from the end of the video. e.g., -1 is last frame" }), @@ -34,7 +34,10 @@ class WanVACEExtend: RETURN_NAMES = ("control_video", "control_mask", "width", "height", "length", "start_images") FUNCTION = "vace_extend" CATEGORY = "video/VACE" - DESCRIPTION = "Generates VACE control video and mask for extending a video from an arbitrary position using context frames." + DESCRIPTION = """ + Generates VACE control video and mask for extending a video from an + arbitrary position using context frames. + """ def vace_extend(self, video, extend_from_idx, context_frames, new_frames): height = int(video.shape[1]) @@ -98,7 +101,6 @@ class WanVACEExtend: return (control_video, mask, width, height, length, start_images) - NODE_CLASS_MAPPINGS = { "WanVACEExtend": WanVACEExtend } diff --git a/wan_vace_prep.py b/wan_vace_prep.py index a487fd4..6d48b0d 100644 --- a/wan_vace_prep.py +++ b/wan_vace_prep.py @@ -1,12 +1,6 @@ import torch - class WanVACEPrepBase: - """ - Shared implementation for WanVACEPrep and WanVACEPrepBatch. - Not registered as a ComfyUI node. - """ - def _validate_dimensions(self, video_1, video_2): """Validate that both videos have matching, 16-divisible dimensions.""" height = int(video_1.shape[1]) @@ -95,7 +89,10 @@ class WanVACEPrep(WanVACEPrepBase): RETURN_NAMES = ("control_video", "control_mask", "width", "height", "length", "start_images", "end_images") FUNCTION = "vace_prep" CATEGORY = "video/VACE" - DESCRIPTION = "Generates VACE control video and mask for smooth transitions between two videos using context frames and frame replacement." + DESCRIPTION = """ + Generates VACE control video and mask for smooth transitions between + two videos using context frames and frame replacement. + """ def vace_prep(self, video_1, video_2, context_frames, replace_frames, new_frames): width, height = self._validate_dimensions(video_1, video_2) @@ -241,14 +238,12 @@ class WanVACEPrepBatch(WanVACEPrepBase): return (control_video, mask, width, height, length, start_images, end_images) - -# ComfyUI node registration NODE_CLASS_MAPPINGS = { "WanVACEPrep": WanVACEPrep, - "WanVACEPrepBatch": WanVACEPrepBatch, + "WanVACEPrepBatch": WanVACEPrepBatch } NODE_DISPLAY_NAME_MAPPINGS = { "WanVACEPrep": "🪐 VACE Join", - "WanVACEPrepBatch": "🪐 VACE Join (Batch)", + "WanVACEPrepBatch": "🪐 VACE Join (Batch)" } diff --git a/wan_vace_prep_inline.py b/wan_vace_prep_inline.py index 24114a4..6c457f3 100644 --- a/wan_vace_prep_inline.py +++ b/wan_vace_prep_inline.py @@ -23,7 +23,7 @@ class VACESmooth(WanVACEPrepBase): end_frame : INT — last frame of the selected range, exclusive (from VFS) context_frames : frames of context to pull from each side of the split replace_frames : frames to regenerate at each transition edge - add_frames : new frames to insert between the two sides + new_frames : new frames to insert between the two sides (equivalent to new_frames in WanVACEPrep) Output signature mirrors WanVACEPrep exactly so the downstream @@ -40,13 +40,13 @@ class VACESmooth(WanVACEPrepBase): "start_frame": ("INT", { "default": 0, "min": 0, - "max": 99999, + "max": 1000000, "tooltip": "First frame of the VFS selection (0-based index)." }), "end_frame": ("INT", { "default": 0, "min": 0, - "max": 99999, + "max": 1000000, "tooltip": "Last frame of the VFS selection, exclusive (0-based index)." }), "context_frames": ("INT", { @@ -63,7 +63,7 @@ class VACESmooth(WanVACEPrepBase): "step": 4, "tooltip": "Number of frames to regenerate at each transition edge (multiple of 4)." }), - "add_frames": ("INT", { + "new_frames": ("INT", { "default": 0, "min": 0, "max": 240, @@ -77,17 +77,16 @@ class VACESmooth(WanVACEPrepBase): RETURN_NAMES = ("control_video", "control_mask", "width", "height", "length", "start_images", "end_images") FUNCTION = "vace_prep_inline" CATEGORY = "video/VACE" - DESCRIPTION = ( - "Inline VACE prep: splits a single video at the VFS selection boundary " - "and generates the control video and mask for smooth in-clip transitions. " - "Output signature matches Wan VACE Prep for drop-in compatibility." - ) + DESCRIPTION = """ + Generates VACE control video and mask for smooth transitions between + two specified frames, replacing and/or adding frames in the transition. + """ # ------------------------------------------------------------------ # Public entry point # ------------------------------------------------------------------ - def vace_prep_inline(self, images, start_frame, end_frame, context_frames, replace_frames, add_frames): + def vace_prep_inline(self, images, start_frame, end_frame, context_frames, replace_frames, new_frames): total_frames = images.shape[0] # ── 1. Validate selection bounds ──────────────────────────────── @@ -129,7 +128,7 @@ class VACESmooth(WanVACEPrepBase): v1_context, v2_context = self._extract_context(video_1, video_2, context_frames, replace_frames) control_video, mask, _ = self._build_control_video_and_mask( video_1, v1_context, v2_context, - context_frames, replace_frames, add_frames, + context_frames, replace_frames, new_frames, height, width ) @@ -174,12 +173,10 @@ class VACESmooth(WanVACEPrepBase): f"The split requires at least context_frames + replace_frames frames on each side." ) - -# ComfyUI node registration NODE_CLASS_MAPPINGS = { - "VACESmooth": VACESmooth, + "VACESmooth": VACESmooth } NODE_DISPLAY_NAME_MAPPINGS = { - "VACESmooth": "🪐 VACE Smooth", + "VACESmooth": "🪐 VACE Smooth" } diff --git a/web/visual_frame_selector.js b/web/visual_frame_selector.js index 5a5824c..ae31799 100644 --- a/web/visual_frame_selector.js +++ b/web/visual_frame_selector.js @@ -155,10 +155,9 @@ function updateStatusLine(ctx) { if (!wasVisible && ctx._remeasureControls) ctx._remeasureControls(); const count = state.selection.endFrame - state.selection.startFrame + 1; - const fps = state.video.fps; const total = state.video.totalFrames; - dom.statusText.textContent = `${fps} FPS | ${total} total | Selected: ${count} frames`; + dom.statusText.textContent = `Total frames: ${total} | Selected frames: ${count}`; } // ── DOM Construction ─────────────────────────────────────────────── @@ -387,11 +386,6 @@ function hookVideoElement(ctx, video) { return true; } -/** - * Re-find and re-hook the native video element after a video file change. - * ComfyUI may create a new