e4db1fc892
- Fixed duration_seconds input - Both duration settings are now shown by default - Audio latent fix
412 lines
21 KiB
Python
412 lines
21 KiB
Python
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) |