scofanoandClaude Haiku 4.5
fe39b72ab2
Add low-risk NVENC encoding option to ImageTransitionNode, VideoImageOverlay, and VideoSplitterNode. Each node now includes a use_gpu parameter (default: False) that enables hevc_nvenc when available, with automatic fallback to libx264 if NVIDIA GPU is unavailable or NVENC fails. Changes: - ImageTransitionNode: Add use_gpu parameter for MP4 frame encoding - VideoImageOverlay: Add use_gpu parameter for overlay video encoding - VideoSplitterNode: Add use_gpu parameter for segment encoding Expected gains: 5-10x speedup for GPU users (hevc_nvenc vs libx264) Risk level: VERY LOW (zero breaking changes, backward compatible) Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
185 lines
7.0 KiB
Python
185 lines
7.0 KiB
Python
import os
|
|
import tempfile
|
|
from pathlib import Path
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw
|
|
import subprocess
|
|
import shutil
|
|
import comfy.utils
|
|
|
|
# ComfyUI folder_paths
|
|
try:
|
|
import folder_paths
|
|
OUTPUT_DIR = folder_paths.get_output_directory()
|
|
except ImportError:
|
|
OUTPUT_DIR = Path("output")
|
|
|
|
FFMPEG = shutil.which("ffmpeg")
|
|
if not FFMPEG:
|
|
raise RuntimeError("ffmpeg not found on PATH.")
|
|
|
|
def tensor_to_pil(tensor):
|
|
"""Convert ComfyUI IMAGE tensor [H,W,C] to PIL Image"""
|
|
arr = (tensor.clamp(0, 1) * 255).byte().numpy()
|
|
return Image.fromarray(arr).convert("RGB")
|
|
|
|
def pil_to_tensor(pil_img):
|
|
"""Convert PIL to ComfyUI IMAGE tensor [H,W,C]"""
|
|
arr = np.array(pil_img).astype(np.float32) / 255.0
|
|
return torch.from_numpy(arr)
|
|
|
|
def hex_to_rgb(hex_color):
|
|
"""Convert hex string to RGB tuple"""
|
|
hex_color = hex_color.lstrip('#')
|
|
return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
|
|
|
|
class ImageTransitionNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image1": ("IMAGE",),
|
|
"image2": ("IMAGE",),
|
|
"duration": ("FLOAT", {"default": 1.0, "min": 0.1, "step": 0.1}),
|
|
"direction": (["Vertical - Down", "Vertical - Up", "Horizontal - Left", "Horizontal - Right"], {"default": "Horizontal - Right"}),
|
|
"line_toggle": ("BOOLEAN", {"default": False}),
|
|
"thickness": ("INT", {"default": 2, "min": 1, "max": 50}),
|
|
"hex_color": ("STRING", {"default": "#FFFFFF"}),
|
|
"prefix": ("STRING", {"default": "transition"}),
|
|
},
|
|
"optional": {
|
|
"use_gpu": ("BOOLEAN", {"default": False, "label": "Use GPU encoding (NVENC) if available"})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("output_path",)
|
|
FUNCTION = "create_transition"
|
|
CATEGORY = "Simple Video Effects/Image Processing"
|
|
OUTPUT_NODE = True
|
|
# OUTPUT_IS_LIST = False
|
|
# OUTPUT_NODE = False
|
|
|
|
OPTIONAL_OUTPUTS = ("output_path",)
|
|
|
|
def create_transition(self, image1, image2, duration, direction, line_toggle, thickness, hex_color, prefix, use_gpu=False, output_path=None):
|
|
# Extract single frames
|
|
if image1.shape[0] != 1 or image2.shape[0] != 1:
|
|
raise ValueError("Inputs must be single images, not batches.")
|
|
img1_pil = tensor_to_pil(image1[0])
|
|
img2_pil = tensor_to_pil(image2[0])
|
|
|
|
# Ensure same size and even dimensions for video encoding
|
|
width, height = img1_pil.size
|
|
# Make dimensions even
|
|
width = width if width % 2 == 0 else width - 1
|
|
height = height if height % 2 == 0 else height - 1
|
|
if img1_pil.size != (width, height):
|
|
img1_pil = img1_pil.resize((width, height))
|
|
if img2_pil.size != (width, height):
|
|
img2_pil = img2_pil.resize((width, height))
|
|
|
|
width, height = img1_pil.size
|
|
fps = 30
|
|
# Generate at higher frame rate for smoother animation
|
|
gen_fps = 60
|
|
total_frames = int(duration * gen_fps)
|
|
if total_frames < 1:
|
|
total_frames = 1
|
|
|
|
# Prepare color
|
|
try:
|
|
line_color = hex_to_rgb(hex_color)
|
|
except:
|
|
line_color = (255, 255, 255)
|
|
|
|
# Output dir
|
|
output_dir = Path(OUTPUT_DIR)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Filename with increment
|
|
if not prefix:
|
|
prefix = "transition"
|
|
counter = 1
|
|
while True:
|
|
filename = f"{prefix}_{counter:03d}.mp4"
|
|
output_path = output_dir / filename
|
|
if not output_path.exists():
|
|
break
|
|
counter += 1
|
|
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
tmpdir = Path(tmpdir)
|
|
frames_dir = tmpdir / "frames"
|
|
frames_dir.mkdir()
|
|
|
|
pbar = comfy.utils.ProgressBar(total_frames)
|
|
for i in range(total_frames):
|
|
progress = i / (total_frames - 1) if total_frames > 1 else 1.0
|
|
|
|
# Create mask: black for img1, white for img2
|
|
mask = Image.new("L", (width, height), 0) # 0 = show img1
|
|
draw_mask = ImageDraw.Draw(mask)
|
|
|
|
if direction == "Horizontal - Right":
|
|
reveal_x = int(progress * width)
|
|
draw_mask.rectangle([0, 0, reveal_x, height], fill=255)
|
|
elif direction == "Horizontal - Left":
|
|
reveal_x = int(progress * width)
|
|
draw_mask.rectangle([width - reveal_x, 0, width, height], fill=255)
|
|
elif direction == "Vertical - Down":
|
|
reveal_y = int(progress * height)
|
|
draw_mask.rectangle([0, 0, width, reveal_y], fill=255)
|
|
elif direction == "Vertical - Up":
|
|
reveal_y = int(progress * height)
|
|
draw_mask.rectangle([0, height - reveal_y, width, height], fill=255)
|
|
|
|
# Composite images
|
|
frame = Image.composite(img2_pil, img1_pil, mask)
|
|
|
|
# Draw line if enabled
|
|
if line_toggle and progress < 1.0:
|
|
draw = ImageDraw.Draw(frame)
|
|
if direction == "Horizontal - Right":
|
|
x = int(progress * (width - thickness))
|
|
draw.line([x, 0, x, height], fill=line_color, width=thickness)
|
|
elif direction == "Horizontal - Left":
|
|
x = int(width - progress * (width - thickness))
|
|
draw.line([x, 0, x, height], fill=line_color, width=thickness)
|
|
elif direction == "Vertical - Down":
|
|
y = int(progress * (height - thickness))
|
|
draw.line([0, y, width, y], fill=line_color, width=thickness)
|
|
elif direction == "Vertical - Up":
|
|
y = int(height - progress * (height - thickness))
|
|
draw.line([0, y, width, y], fill=line_color, width=thickness)
|
|
|
|
# Save frame
|
|
frame.save(frames_dir / f"frame_{i:06d}.png")
|
|
pbar.update(1)
|
|
|
|
# Encode to MP4 (try GPU if requested, fallback to CPU)
|
|
vcodec = 'hevc_nvenc' if use_gpu else 'libx264'
|
|
cmd = [
|
|
FFMPEG,
|
|
"-y",
|
|
"-framerate", str(fps),
|
|
"-i", str(frames_dir / "frame_%06d.png"),
|
|
"-c:v", vcodec,
|
|
"-pix_fmt", "yuv420p",
|
|
str(output_path)
|
|
]
|
|
|
|
try:
|
|
subprocess.run(cmd, check=True)
|
|
except subprocess.CalledProcessError as e:
|
|
# If NVENC fails, fallback to libx264
|
|
if use_gpu and 'hevc_nvenc' in vcodec:
|
|
print(f"NVENC encoding failed, falling back to libx264...")
|
|
cmd[cmd.index(vcodec)] = 'libx264'
|
|
subprocess.run(cmd, check=True)
|
|
else:
|
|
raise
|
|
|
|
return (str(output_path),)
|