Refactor: IMAGE tensor I/O — connect directly to video nodes
- Input/output are now standard ComfyUI IMAGE tensors [B, H, W, C] float32 in [0,1], compatible with any Load/Save Video node - Remove file-path and folder_paths dependencies from the node - Add interpolate_batch_2x() to phase_interpolation.py for frame-list based processing - Output frame count: B*2-1 (original frames interleaved with synthesized midpoints) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
m1_4k
co-authored by
Claude Sonnet 4.6
Claude Sonnet 4.6
parent
1527cfeb90
commit
f0d2ae802a
@@ -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,)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
+33
-1
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user