diff --git a/nodes.py b/nodes.py index d0322ee..8c56a17 100644 --- a/nodes.py +++ b/nodes.py @@ -1,48 +1,39 @@ """ ComfyUI custom node: Phase-Based Frame Interpolation (2x) -Interpolates video frames using Fourier-transform-based phase differences -on a complex steerable pyramid, producing a 2x frame-rate video with -minimal ghosting compared to naive linear interpolation. +Input : IMAGE tensor [B, H, W, C] float32 [0, 1] — batch of video frames +Output : IMAGE tensor [B*2-1, H, W, C] float32 [0, 1] — 2x interpolated frames + +Connects directly to any Load Video / Save Video node that uses the +standard ComfyUI IMAGE batch format. """ -import os import numpy as np -import cv2 -import folder_paths +import torch -from .phase_interpolation import interpolate_video_2x - -VIDEO_EXTENSIONS = ('.mp4', '.avi', '.mov', '.webm', '.mkv', '.gif', '.m4v', '.flv') - - -def _list_input_videos(): - input_dir = folder_paths.get_input_directory() - if not os.path.isdir(input_dir): - return [] - return sorted( - f for f in os.listdir(input_dir) - if f.lower().endswith(VIDEO_EXTENSIONS) - ) +from .phase_interpolation import interpolate_batch_2x class PhaseBasedFrameInterpolation: """ 2x Frame Interpolation via Phase-Based Fourier Method. - Accepts a video file from the ComfyUI input directory, - inserts one interpolated frame between every pair of consecutive frames, - and writes the result to the ComfyUI output directory. + Takes a batch of video frames as a standard ComfyUI IMAGE tensor, + inserts one phase-interpolated frame between every consecutive pair, + and returns the doubled frame batch as an IMAGE tensor. + + Connect: Load Video → [this node] → Save Video / Preview """ @classmethod def INPUT_TYPES(cls): - videos = _list_input_videos() return { "required": { - "video": ( - videos, - {"video_upload": True}, + "images": ( + "IMAGE", + { + "tooltip": "Batch of video frames [B, H, W, C] from a Load Video node.", + }, ), "pyramid_levels": ( "INT", @@ -57,79 +48,40 @@ class PhaseBasedFrameInterpolation: ), }, ), - "output_filename": ( - "STRING", - { - "default": "phase_interp_output.mp4", - "multiline": False, - "tooltip": "Output filename (saved in ComfyUI output directory)", - }, - ), } } - RETURN_TYPES = ("STRING",) - RETURN_NAMES = ("output_video_path",) - OUTPUT_NODE = True + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("images",) FUNCTION = "interpolate" CATEGORY = "video/interpolation" DESCRIPTION = ( "2x frame interpolation using Fourier-transform phase differences. " - "Inserts one synthesized frame between every consecutive pair, " - "doubling the frame rate while preserving sharpness and structure." + "Accepts an IMAGE batch from a Load Video node, inserts one synthesized " + "frame between every consecutive pair, and returns the 2x frame batch. " + "Connect the output to a Save Video or Preview node." ) - def interpolate(self, video: str, pyramid_levels: int, output_filename: str): - # Resolve input path from ComfyUI input directory - input_dir = folder_paths.get_input_directory() - video_path = folder_paths.get_annotated_filepath(video) \ - if "[" in video else os.path.join(input_dir, video) + def interpolate(self, images: torch.Tensor, pyramid_levels: int): + # images: [B, H, W, C] float32 in [0, 1] + B, H, W, C = images.shape + if B < 2: + raise ValueError("Need at least 2 frames for interpolation.") - if not os.path.isfile(video_path): - raise FileNotFoundError(f"Video not found: {video_path}") + # Convert tensor to list of uint8 RGB numpy arrays + np_frames = (images.cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8) + frame_list = [np_frames[i] for i in range(B)] - # Resolve output path - output_dir = folder_paths.get_output_directory() - os.makedirs(output_dir, exist_ok=True) + print(f"[PhaseInterp] Input: {B} frames Output: {B * 2 - 1} frames") - base, ext = os.path.splitext(output_filename.strip() or "phase_interp_output.mp4") - if not ext: - ext = ".mp4" - out_filename = base + ext - out_path = os.path.join(output_dir, out_filename) + # Run phase-based 2x interpolation + out_list = interpolate_batch_2x(frame_list, num_levels=pyramid_levels) - # Run phase-based interpolation - interpolate_video_2x( - input_path=video_path, - output_path=out_path, - num_levels=pyramid_levels, - ) + # Convert back to [B_out, H, W, C] float32 tensor in [0, 1] + out_np = np.stack(out_list, axis=0).astype(np.float32) / 255.0 + out_tensor = torch.from_numpy(out_np) - # Extract first frame as thumbnail preview for the UI - previews = [] - cap = cv2.VideoCapture(out_path) - ret, frame = cap.read() - cap.release() - if ret: - thumb_name = base + "_thumb.png" - thumb_path = os.path.join(output_dir, thumb_name) - cv2.imwrite(thumb_path, frame) - previews = [{"filename": thumb_name, "subfolder": "", "type": "output"}] - - return { - "ui": { - "video": [ - { - "filename": out_filename, - "subfolder": "", - "type": "output", - "format": "video/mp4", - } - ], - "images": previews, - }, - "result": (out_path,), - } + return (out_tensor,) # --------------------------------------------------------------------------- diff --git a/phase_interpolation.py b/phase_interpolation.py index 5b7c816..49db582 100644 --- a/phase_interpolation.py +++ b/phase_interpolation.py @@ -310,7 +310,39 @@ def interpolate_frames(frame1: np.ndarray, frame2: np.ndarray, # --------------------------------------------------------------------------- -# Video processing +# Batch (tensor-based) interpolation — used by the ComfyUI node +# --------------------------------------------------------------------------- + +def interpolate_batch_2x(frames: list, num_levels: int = 6) -> list: + """ + Insert one interpolated frame between every consecutive pair. + + Parameters + ---------- + frames : list of (H, W, 3) uint8 RGB arrays + num_levels : steerable pyramid levels + + Returns + ------- + out_frames : list length 2*N-1, alternating originals and interpolated + """ + if len(frames) < 2: + return frames + + out = [] + for i, frame in enumerate(frames): + out.append(frame) + if i < len(frames) - 1: + interp = interpolate_frames(frame, frames[i + 1], + alpha=0.5, num_levels=num_levels) + out.append(interp) + if (i + 1) % 10 == 0 or i == 0: + print(f"[PhaseInterp] {i + 1}/{len(frames) - 1} pairs done") + return out + + +# --------------------------------------------------------------------------- +# Video file processing (kept for standalone / CLI use) # --------------------------------------------------------------------------- def interpolate_video_2x(input_path: str, output_path: str,