331 lines
13 KiB
Python
331 lines
13 KiB
Python
import os
|
||
import io
|
||
import hashlib
|
||
import folder_paths
|
||
import numpy as np
|
||
import torch
|
||
from PIL import Image
|
||
import cv2
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Helper utilities
|
||
# ---------------------------------------------------------------------------
|
||
|
||
# Safety limits for memory estimation
|
||
DEFAULT_MAX_MEMORY_BYTES = 8 * 1024**3 # hard cap if psutil unavailable
|
||
SYSTEM_RAM_FRACTION_CAP = 0.75 # never exceed 75% of available RAM
|
||
|
||
|
||
def estimate_peak_memory(width: int, height: int,
|
||
num_frames: int, custom_num_frames: int) -> int:
|
||
"""
|
||
Estimate peak RAM in bytes used during frame extraction.
|
||
|
||
Breakdown per call to extract_frames:
|
||
• uint8 staging buffer N × H × W × 3 × 1
|
||
• float32 output tensor N × H × W × 3 × 4
|
||
Both live in memory simultaneously at peak.
|
||
|
||
Plus we keep around:
|
||
• frames (float32, N frames)
|
||
• frames_reversed (float32, same N — torch.flip allocates new storage)
|
||
• frames_at_fps (float32, custom_N frames)
|
||
"""
|
||
bpp_u8 = width * height * 3 # bytes per frame, uint8
|
||
bpp_f32 = bpp_u8 * 4 # bytes per frame, float32
|
||
|
||
# Peak during the first extract_frames() call
|
||
first_call_peak = (num_frames * bpp_u8) + (num_frames * bpp_f32)
|
||
|
||
# After first call completes, frames (f32) + frames_reversed (f32) are resident
|
||
resident_after_first = (num_frames * bpp_f32) * 2
|
||
|
||
# Peak during the second extract_frames() call (target-fps)
|
||
second_call_peak = (
|
||
resident_after_first
|
||
+ (custom_num_frames * bpp_u8)
|
||
+ (custom_num_frames * bpp_f32)
|
||
)
|
||
|
||
return max(first_call_peak, second_call_peak)
|
||
|
||
|
||
def get_memory_limit() -> int:
|
||
"""Return the max allowed peak memory in bytes."""
|
||
try:
|
||
import psutil
|
||
available = psutil.virtual_memory().available
|
||
return int(available * SYSTEM_RAM_FRACTION_CAP)
|
||
except Exception:
|
||
return DEFAULT_MAX_MEMORY_BYTES
|
||
|
||
|
||
def format_bytes(n: int) -> str:
|
||
for unit in ("B", "KB", "MB", "GB", "TB"):
|
||
if abs(n) < 1024.0:
|
||
return f"{n:.1f} {unit}"
|
||
n /= 1024.0
|
||
return f"{n:.1f} PB"
|
||
|
||
|
||
def get_video_info(video_path: str) -> dict:
|
||
"""Return basic metadata for a video file."""
|
||
cap = cv2.VideoCapture(video_path)
|
||
if not cap.isOpened():
|
||
raise ValueError(f"Cannot open video: {video_path}")
|
||
|
||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||
total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||
cap.release()
|
||
|
||
return {
|
||
"fps": fps,
|
||
"total_frames": total,
|
||
"width": width,
|
||
"height": height,
|
||
"duration_seconds": total / fps if fps > 0 else 0,
|
||
"bytes_per_frame_f32": width * height * 3 * 4,
|
||
"memory_limit_bytes": get_memory_limit(),
|
||
}
|
||
|
||
|
||
def read_single_frame(cap: cv2.VideoCapture, frame_index: int,
|
||
width: int, height: int) -> np.ndarray:
|
||
"""Read one frame; return a black frame on failure."""
|
||
cap.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
|
||
ret, frame = cap.read()
|
||
if not ret:
|
||
return np.zeros((height, width, 3), dtype=np.uint8)
|
||
return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||
|
||
|
||
def extract_frames(video_path: str, start_frame: int, end_frame: int,
|
||
num_frames: int) -> torch.Tensor:
|
||
"""
|
||
Extract *num_frames* evenly-spaced frames from [start_frame, end_frame].
|
||
Returns a BHWC float32 tensor in [0, 1].
|
||
"""
|
||
cap = cv2.VideoCapture(video_path)
|
||
if not cap.isOpened():
|
||
raise ValueError(f"Cannot open video: {video_path}")
|
||
|
||
total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||
start_frame = max(0, min(start_frame, total - 1))
|
||
end_frame = max(start_frame, min(end_frame, total - 1))
|
||
num_frames = max(1, num_frames)
|
||
|
||
if start_frame == end_frame:
|
||
indices = [start_frame]
|
||
else:
|
||
indices = [
|
||
int(round(start_frame + i * (end_frame - start_frame) / (num_frames - 1)))
|
||
for i in range(num_frames)
|
||
]
|
||
|
||
arr = np.empty((len(indices), h, w, 3), dtype=np.uint8)
|
||
for i, idx in enumerate(indices):
|
||
arr[i] = read_single_frame(cap, idx, w, h)
|
||
cap.release()
|
||
|
||
float_arr = arr.astype(np.float32) / 255.0
|
||
del arr
|
||
return torch.from_numpy(float_arr) # BHWC
|
||
|
||
|
||
def frame_as_tensor(video_path: str, frame_index: int) -> torch.Tensor:
|
||
"""Return a single frame as a (1, H, W, 3) float32 tensor."""
|
||
cap = cv2.VideoCapture(video_path)
|
||
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||
frame = read_single_frame(cap, frame_index, w, h)
|
||
cap.release()
|
||
arr = frame.astype(np.float32) / 255.0
|
||
return torch.from_numpy(arr).unsqueeze(0) # 1HWC
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Node definition
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class VideoFrameExtractor:
|
||
"""
|
||
ComfyUI node: VideoFrameExtractor
|
||
─────────────────────────────────
|
||
Outputs based on the loop region selected in the timeline widget:
|
||
• first_frame — single IMAGE of the loop start frame
|
||
• last_frame — single IMAGE of the loop end frame
|
||
• frames — all extracted frames in forward order (BHWC)
|
||
• frames_reversed — same frames in reverse order (BHWC)
|
||
• filename — video filename (STRING)
|
||
• width — video width in pixels (INT)
|
||
• height — video height in pixels (INT)
|
||
• fps — source frame rate (FLOAT)
|
||
"""
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
video_files = []
|
||
input_dir = folder_paths.get_input_directory()
|
||
for fname in sorted(os.listdir(input_dir)):
|
||
if fname.lower().endswith((".mp4", ".avi", ".mov", ".mkv", ".webm")):
|
||
video_files.append(fname)
|
||
if not video_files:
|
||
video_files = [""]
|
||
|
||
return {
|
||
"required": {
|
||
"video": (video_files, {
|
||
"video_upload": True,
|
||
}),
|
||
"start_frame": ("INT", {
|
||
"default": 0, "min": 0, "max": 999999, "step": 1,
|
||
"display": "number",
|
||
}),
|
||
"end_frame": ("INT", {
|
||
"default": 0, "min": 0, "max": 999999, "step": 1,
|
||
"display": "number",
|
||
}),
|
||
"num_frames": ("INT", {
|
||
"default": 0, "min": 0, "max": 999999, "step": 1,
|
||
"display": "number",
|
||
}),
|
||
"target_fps": ("FLOAT", {
|
||
"default": 24.0, "min": 0.1, "max": 240.0, "step": 0.1,
|
||
"display": "number",
|
||
}),
|
||
},
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE", "STRING", "INT", "INT", "FLOAT", "IMAGE", "FLOAT")
|
||
RETURN_NAMES = ("Clipped Frames", "Reversed Clipped Frames", "First Frame", "Last Frame",
|
||
"Filename Prefix (no extension)", "Width", "Height", "Original FPS",
|
||
"Clipped Frames at Target FPS", "Target FPS")
|
||
FUNCTION = "extract"
|
||
CATEGORY = "video"
|
||
OUTPUT_NODE = True
|
||
|
||
@classmethod
|
||
def IS_CHANGED(cls, video, start_frame, end_frame, num_frames, target_fps):
|
||
video_path = folder_paths.get_annotated_filepath(video)
|
||
if not os.path.exists(video_path):
|
||
return float("nan")
|
||
m = hashlib.md5()
|
||
m.update(f"{video}{start_frame}{end_frame}{num_frames}{target_fps}".encode())
|
||
return m.hexdigest()
|
||
|
||
def extract(self, video: str, start_frame: int, end_frame: int,
|
||
num_frames: int, target_fps: float):
|
||
video_path = folder_paths.get_annotated_filepath(video)
|
||
if not os.path.exists(video_path):
|
||
raise FileNotFoundError(f"Video not found: {video_path}")
|
||
|
||
info = get_video_info(video_path)
|
||
total = info["total_frames"]
|
||
width = info["width"]
|
||
height = info["height"]
|
||
fps = info["fps"]
|
||
|
||
# Clamp loop region
|
||
start_frame = max(0, min(start_frame, total - 1))
|
||
end_frame = max(start_frame, min(end_frame, total - 1))
|
||
if end_frame == 0 and start_frame == 0:
|
||
end_frame = total - 1
|
||
|
||
# num_frames == 0 → auto (JS widget sets this to loop span)
|
||
if num_frames <= 0:
|
||
num_frames = end_frame - start_frame + 1
|
||
|
||
# ── Memory safety check ──────────────────────────────────────────
|
||
clip_duration = (end_frame - start_frame + 1) / fps if fps > 0 else 0
|
||
custom_num_frames = max(1, round(clip_duration * target_fps))
|
||
peak_bytes = estimate_peak_memory(width, height, num_frames,
|
||
custom_num_frames)
|
||
limit_bytes = get_memory_limit()
|
||
|
||
if peak_bytes > limit_bytes:
|
||
raise RuntimeError(
|
||
f"Refusing to run: estimated peak memory "
|
||
f"{format_bytes(peak_bytes)} exceeds the safe limit of "
|
||
f"{format_bytes(limit_bytes)}.\n\n"
|
||
f"Clip: {num_frames} frames @ {width}x{height} "
|
||
f"(+{custom_num_frames} frames at {target_fps:.1f} fps).\n\n"
|
||
f"Reduce the loop region, shrink the video resolution, or "
|
||
f"lower target_fps. A shorter clip uses proportionally less RAM."
|
||
)
|
||
|
||
# Extract the frame batch
|
||
frames = extract_frames(video_path, start_frame, end_frame, num_frames)
|
||
frames_reversed = torch.flip(frames, dims=[0])
|
||
|
||
# Single-frame outputs for the loop boundaries — slice from already-loaded batch
|
||
first_frame = frames[0:1]
|
||
last_frame = frames[-1:]
|
||
|
||
# Filename prefix — basename without extension
|
||
filename_prefix = os.path.splitext(os.path.basename(video_path))[0]
|
||
|
||
# Custom FPS batch: resample the loop region to target_fps
|
||
# (custom_num_frames was computed above for the memory check)
|
||
frames_at_fps = extract_frames(video_path, start_frame, end_frame, custom_num_frames)
|
||
|
||
return (frames, frames_reversed, first_frame, last_frame,
|
||
filename_prefix, width, height, fps, frames_at_fps, float(target_fps))
|
||
|
||
|
||
# ─── API routes ──────────────────────────────────────────────────────────────
|
||
try:
|
||
from server import PromptServer
|
||
from aiohttp import web
|
||
|
||
@PromptServer.instance.routes.get("/video_frame_extractor/info")
|
||
async def video_info_handler(request):
|
||
video_name = request.rel_url.query.get("video", "")
|
||
if not video_name:
|
||
return web.json_response({"error": "no video specified"}, status=400)
|
||
try:
|
||
video_path = folder_paths.get_annotated_filepath(video_name)
|
||
info = get_video_info(video_path)
|
||
return web.json_response(info)
|
||
except Exception as exc:
|
||
return web.json_response({"error": str(exc)}, status=500)
|
||
|
||
@PromptServer.instance.routes.get("/video_frame_extractor/thumbnail")
|
||
async def video_thumbnail_handler(request):
|
||
"""Return a single JPEG frame at the requested position."""
|
||
video_name = request.rel_url.query.get("video", "")
|
||
frame_index = int(request.rel_url.query.get("frame", 0))
|
||
if not video_name:
|
||
return web.Response(status=400)
|
||
try:
|
||
video_path = folder_paths.get_annotated_filepath(video_name)
|
||
cap = cv2.VideoCapture(video_path)
|
||
cap.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
|
||
ret, frame = cap.read()
|
||
cap.release()
|
||
if not ret:
|
||
return web.Response(status=404)
|
||
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||
img = Image.fromarray(frame_rgb)
|
||
img.thumbnail((640, 360))
|
||
buf = io.BytesIO()
|
||
img.save(buf, format="JPEG", quality=80)
|
||
body = buf.getvalue()
|
||
buf.close()
|
||
return web.Response(body=body, content_type="image/jpeg")
|
||
except Exception as exc:
|
||
return web.Response(status=500, text=str(exc))
|
||
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
# ─── Exports ─────────────────────────────────────────────────────────────────
|
||
|
||
NODE_CLASS_MAPPINGS = {"VideoFrameExtractor": VideoFrameExtractor}
|
||
NODE_DISPLAY_NAME_MAPPINGS = {"VideoFrameExtractor": "Video Frame Extractor 🎞️"}
|