diff --git a/load_video_ui.py b/load_video_ui.py deleted file mode 100644 index fc1801c..0000000 --- a/load_video_ui.py +++ /dev/null @@ -1,412 +0,0 @@ -import os -import torch -import numpy as np -import folder_paths -import av -from server import PromptServer -from aiohttp import web -import comfy.utils - -# Custom API route to serve video files from anywhere on the user's system for the frontend preview -@PromptServer.instance.routes.get("/video_ui_custom_view") -async def custom_view(request): - file_path = request.query.get("filename", "") - if os.path.exists(file_path) and os.path.isfile(file_path): - return web.FileResponse(file_path) - return web.Response(status=404, text="File not found") - -# Custom API route for Chunked Uploads to bypass the 413 Payload Too Large error -@PromptServer.instance.routes.post("/video_ui_upload_chunk") -async def upload_chunk(request): - post = await request.post() - file = post.get("file") - filename = post.get("filename") - chunk_index = int(post.get("chunk_index")) - total_chunks = int(post.get("total_chunks")) - - upload_dir = folder_paths.get_input_directory() - file_path = os.path.join(upload_dir, filename) - - # Append to file if it's not the first chunk, otherwise write new - mode = "ab" if chunk_index > 0 else "wb" - with open(file_path, mode) as f: - f.write(file.file.read()) - - if chunk_index == total_chunks - 1: - return web.json_response({"name": filename}) - return web.json_response({"status": "ok"}) - - -class LoadVideoUI: - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "video": ("STRING", {"default": ""}), - "start_time": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100000.0, "step": 0.01}), - "end_time": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100000.0, "step": 0.01}), - "duration": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 100000.0, "step": 0.01}), - "start_frame": ("INT", {"default": 0, "min": 0, "max": 10000000, "step": 1}), - "end_frame": ("INT", {"default": 0, "min": 0, "max": 10000000, "step": 1}), - "duration_frames": ("INT", {"default": 0, "min": 0, "max": 10000000, "step": 1}), - "resize_method": (["maintain aspect ratio", "stretch to fit", "pad", "crop"], {"default": "maintain aspect ratio"}), - "custom_width": ("INT", {"default": 0, "min": 0, "max": 100000, "step": 8, "tooltip": "Custom width. 0 means original width."}), - "custom_height": ("INT", {"default": 0, "min": 0, "max": 100000, "step": 8, "tooltip": "Custom height. 0 means original height."}), - "frame_rate": ("INT", {"default": 24, "min": 1, "max": 120, "step": 1, "tooltip": "Force the video to a specific frame rate for extraction."}), - "display_mode": (["seconds", "frames"], {"default": "seconds"}), - "crop_x": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}), - "crop_y": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}), - "crop_w": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}), - "crop_h": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}), - } - } - - RETURN_TYPES = ("IMAGE", "AUDIO", "FLOAT", "INT") - RETURN_NAMES = ("images", "audio", "duration", "frame_count") - FUNCTION = "load_video" - CATEGORY = "WhatDreamsCost" - - def load_video(self, video, frame_rate, display_mode, start_time, end_time, duration, start_frame, end_frame, duration_frames, custom_width=0, custom_height=0, resize_method="maintain aspect ratio", crop_x=0.0, crop_y=0.0, crop_w=1.0, crop_h=1.0, **kwargs): - if not video: - # Return blank defaults if no video is loaded - empty_image = torch.zeros((1, 512, 512, 3), dtype=torch.float32) - empty_audio = {"waveform": torch.zeros((1, 1, 44100)), "sample_rate": 44100} - return (empty_image, empty_audio, 0.0, 0) - - # 1. Resolve path using ComfyUI standard paths or Absolute Path - video_path = video # Try exact/absolute path first - if not os.path.exists(video_path): - video_path_annotated = folder_paths.get_annotated_filepath(video) - if os.path.exists(video_path_annotated): - video_path = video_path_annotated - else: - video_path_input = os.path.join(folder_paths.get_input_directory(), video) - if os.path.exists(video_path_input): - video_path = video_path_input - else: - raise FileNotFoundError(f"Video file not found: {video}") - - # Open container to read streams and metadata - container = av.open(video_path) - - # Determine video stream and duration - video_stream = container.streams.video[0] if len(container.streams.video) > 0 else None - video_duration = 0 - if video_stream and video_stream.duration and video_stream.time_base: - video_duration = float(video_stream.duration * video_stream.time_base) - - orig_w = video_stream.codec_context.width if video_stream else 512 - orig_h = video_stream.codec_context.height if video_stream else 512 - - # Determine correct colorspace and color range for PyAV conversion to prevent color shift - try: - from av.video.reformatter import Colorspace, ColorRange - # Improve fallback heuristic to check both dimensions (e.g. 720x1280 vertical video is HD) - fallback_cs = Colorspace.ITU709 if max(orig_w, orig_h) >= 720 else Colorspace.ITU601 - fallback_cr = ColorRange.MPEG - dst_range = ColorRange.JPEG # RGB should always be full range - except ImportError: - fallback_cs = "itu709" if max(orig_w, orig_h) >= 720 else "itu601" - fallback_cr = "mpeg" - dst_range = "jpeg" - - src_colorspace = fallback_cs - src_color_range = fallback_cr - - if video_stream and video_stream.codec_context: - cc = video_stream.codec_context - - c_space = getattr(cc, 'colorspace', getattr(cc, 'color_space', None)) - if c_space and hasattr(c_space, 'name') and c_space.name != "UNSPECIFIED": - src_colorspace = c_space - elif c_space and isinstance(c_space, str) and "unspecified" not in c_space.lower(): - src_colorspace = c_space - - c_range = getattr(cc, 'color_range', None) - if c_range and hasattr(c_range, 'name') and c_range.name != "UNSPECIFIED": - src_color_range = c_range - elif c_range and isinstance(c_range, str) and "unspecified" not in c_range.lower(): - src_color_range = c_range - - target_w = custom_width if custom_width > 0 else orig_w - target_h = custom_height if custom_height > 0 else orig_h - - target_w = target_w - (target_w % 2) - target_h = target_h - (target_h % 2) - - # Calculate manual crop from interactive UI first - manual_crop_left = int(orig_w * crop_x) - manual_crop_top = int(orig_h * crop_y) - manual_crop_right = orig_w - int(orig_w * (crop_x + crop_w)) - manual_crop_bottom = orig_h - int(orig_h * (crop_y + crop_h)) - - # Ensure we don't crop more than the image - manual_crop_left = max(0, min(manual_crop_left, orig_w - 1)) - manual_crop_top = max(0, min(manual_crop_top, orig_h - 1)) - manual_crop_right = max(0, min(manual_crop_right, orig_w - manual_crop_left - 1)) - manual_crop_bottom = max(0, min(manual_crop_bottom, orig_h - manual_crop_top - 1)) - - # After manual crop, the new original dimensions are: - cropped_orig_w = orig_w - manual_crop_left - manual_crop_right - cropped_orig_h = orig_h - manual_crop_top - manual_crop_bottom - - # If no custom width/height is provided, use the cropped original dimensions - if custom_width == 0: - target_w = cropped_orig_w - target_w = target_w - (target_w % 2) - if custom_height == 0: - target_h = cropped_orig_h - target_h = target_h - (target_h % 2) - - scale_w, scale_h = target_w, target_h - pad_left = pad_right = pad_top = pad_bottom = 0 - crop_left = crop_right = crop_top = crop_bottom = 0 - - if custom_width > 0 or custom_height > 0: - if resize_method == "maintain aspect ratio" or resize_method == "pad": - ratio = min(target_w / cropped_orig_w, target_h / cropped_orig_h) - scale_w = int(cropped_orig_w * ratio) - scale_h = int(cropped_orig_h * ratio) - scale_w = scale_w - (scale_w % 2) - scale_h = scale_h - (scale_h % 2) - - if resize_method == "pad": - pad_x = target_w - scale_w - pad_y = target_h - scale_h - pad_left = pad_x // 2 - pad_right = pad_x - pad_left - pad_top = pad_y // 2 - pad_bottom = pad_y - pad_top - else: - target_w, target_h = scale_w, scale_h - - elif resize_method == "crop": - ratio = max(target_w / cropped_orig_w, target_h / cropped_orig_h) - scale_w = int(cropped_orig_w * ratio) - scale_h = int(cropped_orig_h * ratio) - scale_w = scale_w - (scale_w % 2) - scale_h = scale_h - (scale_h % 2) - - crop_x = scale_w - target_w - crop_y = scale_h - target_h - crop_left = crop_x // 2 - crop_right = crop_x - crop_left - crop_top = crop_y // 2 - crop_bottom = crop_y - crop_top - - elif resize_method == "stretch to fit": - scale_w, scale_h = target_w, target_h - - # Determine exact bounds based on frontend mode - if display_mode == "frames": - fr = float(frame_rate) if frame_rate > 0 else 24.0 - actual_start_time = float(start_frame) / fr - actual_end_time = float(end_frame) / fr if (end_frame > 0 and end_frame > start_frame) else video_duration - else: - actual_start_time = start_time - actual_end_time = end_time if (end_time > 0 and end_time > start_time) else video_duration - - if actual_end_time <= 0: - actual_end_time = float('inf') # Fallback if duration is unknown - - # 2. Extract Video Frames (PyAV) - frames = [] - image_tensor = None - frames_loaded = 0 - - if video_stream: - video_stream.thread_type = "AUTO" # Enable multithreaded decoding - - # Efficiently seek backwards to the nearest keyframe - if video_stream.time_base: - seek_pts = int(actual_start_time / float(video_stream.time_base)) - else: - seek_pts = int(actual_start_time * av.time_base) - - container.seek(seek_pts, stream=video_stream, backward=True) - - # Custom sampling to force specific framerate - frame_interval = 1.0 / float(frame_rate) if frame_rate > 0 else 1.0/24.0 - expected_target_time = actual_start_time - - # Pre-calculate expected frames - alloc_end_time = actual_end_time if actual_end_time != float('inf') else video_duration - expected_frames = 0 - if alloc_end_time > 0: - duration_to_extract = alloc_end_time - actual_start_time - if duration_to_extract > 0: - expected_frames = int(np.ceil(duration_to_extract / frame_interval)) + 2 - - pbar = comfy.utils.ProgressBar(expected_frames) if expected_frames > 0 else None - - for frame in container.decode(video_stream): - frame_time = frame.time - if frame_time is None: - frame_time = float(frame.pts * float(video_stream.time_base)) if frame.pts and video_stream.time_base else 0.0 - - if frame_time < actual_start_time: - continue - - # Add a slight buffer (interval) to ensure we evaluate the boundary correctly - if frame_time > actual_end_time + frame_interval: - break - - # Fix PyAV color shift by forcing proper colorspace and range conversion. - # Omit dst_colorspace so swscale defaults naturally for RGB output - # (passing it can cause the YUV matrix to be applied incorrectly). - try: - frame = frame.reformat( - format="rgb24", - src_colorspace=src_colorspace, - src_color_range=src_color_range, - dst_color_range=dst_range - ) - frame_rgb = frame.to_ndarray(format='rgb24') - except Exception as e: - # Fallback: if explicit color reformat fails, use PyAV's default conversion - print(f"[LoadVideoUI] Color reformat failed, using default: {e}") - frame_rgb = frame.to_ndarray(format='rgb24') - - # Apply interactive crop first - if manual_crop_left > 0 or manual_crop_top > 0 or manual_crop_right > 0 or manual_crop_bottom > 0: - frame_rgb = frame_rgb[manual_crop_top:orig_h-manual_crop_bottom, manual_crop_left:orig_w-manual_crop_right, :] - - # Now resize to the scaled dimensions - if scale_w != cropped_orig_w or scale_h != cropped_orig_h: - import cv2 - frame_rgb = cv2.resize(frame_rgb, (scale_w, scale_h), interpolation=cv2.INTER_AREA) - - if crop_left > 0 or crop_top > 0 or crop_right > 0 or crop_bottom > 0: - frame_rgb = frame_rgb[crop_top:scale_h-crop_bottom, crop_left:scale_w-crop_right, :] - if pad_left > 0 or pad_top > 0 or pad_right > 0 or pad_bottom > 0: - frame_rgb = np.pad(frame_rgb, ((pad_top, pad_bottom), (pad_left, pad_right), (0, 0)), mode='constant', constant_values=0) - - # Duplicate or skip frames perfectly based on timestamps to meet forced framerate. - # FIX: Use strictly less than (<) for actual_end_time to prevent the loop from fetching an extra +1 frame - # at the exact boundary of the duration slice! - while expected_target_time <= frame_time and expected_target_time < actual_end_time - 1e-5: - if image_tensor is None and expected_frames > 0: - # First frame: allocate the tensor - height, width = frame_rgb.shape[:2] - alloc_frames = expected_frames + 50 # Add generous buffer to prevent reallocation - try: - image_tensor = torch.zeros((alloc_frames, height, width, 3), dtype=torch.float32) - except Exception as e: - print(f"[LoadVideoUI] Pre-allocation failed, falling back to list: {e}") - expected_frames = 0 # Disable pre-allocation - - if image_tensor is not None: - # Check bounds (just in case) - if frames_loaded >= image_tensor.shape[0]: - # Extend tensor if we underestimated - extension = torch.zeros((50, image_tensor.shape[1], image_tensor.shape[2], 3), dtype=torch.float32) - image_tensor = torch.cat((image_tensor, extension), dim=0) - - # Insert frame with minimal memory copy directly to tensor - image_tensor[frames_loaded] = torch.from_numpy(frame_rgb).float().div_(255.0) - frames_loaded += 1 - else: - # Fallback list append if pre-allocation failed - frames.append(frame_rgb) - - if pbar: - pbar.update(1) - - expected_target_time += frame_interval - - # Convert frames to ComfyUI Image standard format [N, H, W, C], float32, range 0.0-1.0 - if image_tensor is not None: - if frames_loaded > 0: - image_tensor = image_tensor[:frames_loaded] - else: - image_tensor = torch.zeros((1, 512, 512, 3), dtype=torch.float32) - elif len(frames) > 0: - frames_np = np.array(frames, dtype=np.float32) / 255.0 - image_tensor = torch.from_numpy(frames_np) - else: - # Fallback for an empty slice - image_tensor = torch.zeros((1, 512, 512, 3), dtype=torch.float32) - - # 3. Extract Audio (PyAV) - audio_dict = {"waveform": torch.zeros((1, 1, 44100)), "sample_rate": 44100} # Default empty audio - - if len(container.streams.audio) > 0: - try: - audio_stream = container.streams.audio[0] - audio_stream.thread_type = "AUTO" - sample_rate = getattr(audio_stream, 'rate', 44100) or 44100 - - # We must seek again on the container specifically for the audio stream - if audio_stream.time_base: - seek_pts = int(actual_start_time / float(audio_stream.time_base)) - else: - seek_pts = int(actual_start_time * av.time_base) - - container.seek(seek_pts, stream=audio_stream, backward=True) - - # Resample to standard float planar format (fltp) - resampler = av.AudioResampler(format='fltp') - - audio_data = [] - first_frame_time = None - - for frame in container.decode(audio_stream): - frame_time = frame.time - if frame_time is None: - frame_time = float(frame.pts * float(audio_stream.time_base)) if frame.pts and audio_stream.time_base else 0.0 - - # Give a small 1-second buffer to ensure we catch end frames - if frame_time > actual_end_time + 1.0: - break - - if first_frame_time is None: - first_frame_time = frame_time - - resampled_frames = resampler.resample(frame) - for r_frame in resampled_frames: - audio_data.append(r_frame.to_ndarray()) - - if audio_data: - # Concatenate all frames horizontally along the sample axis - waveform_np = np.concatenate(audio_data, axis=1) - waveform = torch.from_numpy(waveform_np).float() - - if first_frame_time is None: - first_frame_time = 0.0 - - # Calculate exact slice points to trim precisely - offset_sec = max(0.0, actual_start_time - first_frame_time) - start_sample = int(offset_sec * sample_rate) - - duration_sec_audio = actual_end_time - actual_start_time - end_sample = start_sample + int(duration_sec_audio * sample_rate) - - # Trim array bounds properly - if end_sample > start_sample: - waveform = waveform[:, start_sample:end_sample] - else: - waveform = waveform[:, start_sample:] - - # Expand to ComfyUI Audio standard [batch_size, channels, samples] - waveform = waveform.unsqueeze(0) - audio_dict = {"waveform": waveform, "sample_rate": sample_rate} - - except Exception as e: - # Catch gracefully without breaking the pipeline execution - print(f"[LoadVideoUI] Audio track extraction skipped or failed: {e}") - - # Always close container to free up system memory lock - container.close() - - # Output accurate final duration in seconds - final_duration_sec = float(max(0.0, actual_end_time - actual_start_time)) - - # Accurately output the true number of extracted frames - # (Using the shape of the array provides exact 1:1 parity with the timeline's math) - frame_count = image_tensor.shape[0] if (frames_loaded > 0 or len(frames) > 0) else 0 - if frame_count == 0 and final_duration_sec > 0: - # Fallback estimation only if PyAV completely failed to decode a valid chunk - calc_fr = float(frame_rate) if frame_rate > 0 else 24.0 - frame_count = int(np.floor(final_duration_sec * calc_fr)) - - return (image_tensor, audio_dict, final_duration_sec, frame_count) \ No newline at end of file