292 lines
10 KiB
Python
292 lines
10 KiB
Python
"""
|
|
Utility functions for Frame Utility nodes.
|
|
|
|
Provides shared functionality for frame manipulation, validation, and processing.
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from typing import Tuple, Optional, Union
|
|
import logging
|
|
|
|
def validate_image_tensor(tensor: torch.Tensor, name: str) -> None:
|
|
"""
|
|
Validate that a tensor is a proper IMAGE tensor for ComfyUI.
|
|
|
|
Args:
|
|
tensor: Input tensor to validate
|
|
name: Name of the tensor for error reporting
|
|
|
|
Raises:
|
|
ValueError: If tensor format is invalid
|
|
"""
|
|
if not isinstance(tensor, torch.Tensor):
|
|
raise ValueError(f"{name} must be a torch.Tensor, got {type(tensor)}")
|
|
|
|
if len(tensor.shape) != 4:
|
|
raise ValueError(f"{name} must be 4D tensor [B, H, W, C], got shape {tensor.shape}")
|
|
|
|
if tensor.shape[-1] not in [1, 3, 4]: # Grayscale, RGB, or RGBA
|
|
raise ValueError(f"{name} must have 1, 3, or 4 channels, got {tensor.shape[-1]}")
|
|
|
|
if tensor.dtype not in [torch.float32, torch.float16]:
|
|
logging.warning(f"{name} has dtype {tensor.dtype}, converting to float32")
|
|
|
|
def resize_to_match(source: torch.Tensor, target: torch.Tensor,
|
|
method: str = "bilinear") -> torch.Tensor:
|
|
"""
|
|
Resize source tensor to match target tensor dimensions.
|
|
|
|
Args:
|
|
source: Source tensor to resize [B, H, W, C]
|
|
target: Target tensor to match dimensions [B, H, W, C]
|
|
method: Interpolation method ('nearest', 'bilinear', 'bicubic', 'area')
|
|
|
|
Returns:
|
|
Resized source tensor matching target dimensions
|
|
"""
|
|
if source.shape[1:3] == target.shape[1:3]:
|
|
return source
|
|
|
|
# Convert to [B, C, H, W] for F.interpolate
|
|
source_reshaped = source.permute(0, 3, 1, 2)
|
|
|
|
# Resize to target dimensions
|
|
resized = F.interpolate(
|
|
source_reshaped,
|
|
size=(target.shape[1], target.shape[2]),
|
|
mode=method,
|
|
align_corners=False if method in ['bilinear', 'bicubic'] else None
|
|
)
|
|
|
|
# Convert back to [B, H, W, C]
|
|
return resized.permute(0, 2, 3, 1)
|
|
|
|
def advanced_resize(frames: torch.Tensor, width: int, height: int,
|
|
interpolation: str = "bilinear", method: str = "keep proportion",
|
|
condition: str = "always", multiple_of: int = 0) -> torch.Tensor:
|
|
"""
|
|
Advanced frame resizing with multiple methods and conditions.
|
|
|
|
Args:
|
|
frames: Input frames [B, H, W, C]
|
|
width: Target width
|
|
height: Target height
|
|
interpolation: Interpolation method
|
|
method: Resizing method ('stretch', 'keep proportion', 'fill / crop', 'pad')
|
|
condition: When to resize ('always', 'downscale if bigger', 'upscale if smaller', 'if different')
|
|
multiple_of: Round dimensions to multiple of this value
|
|
|
|
Returns:
|
|
Resized frames
|
|
"""
|
|
current_height, current_width = frames.shape[1], frames.shape[2]
|
|
|
|
# Apply multiple_of constraint
|
|
if multiple_of > 0:
|
|
width = ((width + multiple_of - 1) // multiple_of) * multiple_of
|
|
height = ((height + multiple_of - 1) // multiple_of) * multiple_of
|
|
|
|
# Check condition
|
|
should_resize = True
|
|
if condition == "downscale if bigger":
|
|
should_resize = current_width > width or current_height > height
|
|
elif condition == "upscale if smaller":
|
|
should_resize = current_width < width or current_height < height
|
|
elif condition == "if different":
|
|
should_resize = current_width != width or current_height != height
|
|
|
|
if not should_resize:
|
|
return frames
|
|
|
|
# Calculate target dimensions based on method
|
|
if method == "stretch":
|
|
target_width, target_height = width, height
|
|
elif method == "keep proportion":
|
|
# Maintain aspect ratio, fit within target dimensions
|
|
scale = min(width / current_width, height / current_height)
|
|
target_width = int(current_width * scale)
|
|
target_height = int(current_height * scale)
|
|
|
|
# Apply multiple_of constraint to final dimensions
|
|
if multiple_of > 0:
|
|
target_width = ((target_width + multiple_of - 1) // multiple_of) * multiple_of
|
|
target_height = ((target_height + multiple_of - 1) // multiple_of) * multiple_of
|
|
|
|
elif method == "fill / crop":
|
|
# Fill target dimensions, crop if necessary
|
|
scale = max(width / current_width, height / current_height)
|
|
target_width = int(current_width * scale)
|
|
target_height = int(current_height * scale)
|
|
elif method == "pad":
|
|
# Keep original size, will pad later
|
|
target_width, target_height = current_width, current_height
|
|
else:
|
|
target_width, target_height = width, height
|
|
|
|
# Perform resize
|
|
frames_reshaped = frames.permute(0, 3, 1, 2) # [B, C, H, W]
|
|
|
|
# Handle different interpolation methods
|
|
mode = interpolation
|
|
if interpolation == "lanczos":
|
|
# PyTorch doesn't have lanczos, use bicubic as fallback
|
|
mode = "bicubic"
|
|
|
|
resized = F.interpolate(
|
|
frames_reshaped,
|
|
size=(target_height, target_width),
|
|
mode=mode,
|
|
align_corners=False if mode in ['bilinear', 'bicubic'] else None
|
|
)
|
|
|
|
resized = resized.permute(0, 2, 3, 1) # [B, H, W, C]
|
|
|
|
# Handle padding or cropping for specific methods
|
|
if method == "pad" and (target_width != width or target_height != height):
|
|
# Pad to target dimensions
|
|
pad_h = max(0, height - target_height)
|
|
pad_w = max(0, width - target_width)
|
|
|
|
if pad_h > 0 or pad_w > 0:
|
|
pad_top = pad_h // 2
|
|
pad_bottom = pad_h - pad_top
|
|
pad_left = pad_w // 2
|
|
pad_right = pad_w - pad_left
|
|
|
|
# Pad with zeros (black)
|
|
resized = F.pad(resized.permute(0, 3, 1, 2),
|
|
(pad_left, pad_right, pad_top, pad_bottom),
|
|
mode='constant', value=0.0).permute(0, 2, 3, 1)
|
|
|
|
elif method == "fill / crop" and (target_width != width or target_height != height):
|
|
# Crop to target dimensions
|
|
if target_height > height or target_width > width:
|
|
crop_h = min(target_height, height)
|
|
crop_w = min(target_width, width)
|
|
|
|
start_h = max(0, (target_height - crop_h) // 2)
|
|
start_w = max(0, (target_width - crop_w) // 2)
|
|
|
|
end_h = min(target_height, start_h + crop_h)
|
|
end_w = min(target_width, start_w + crop_w)
|
|
|
|
resized = resized[:, start_h:end_h, start_w:end_w, :]
|
|
|
|
return resized
|
|
|
|
def blend_frames(frame1: torch.Tensor, frame2: torch.Tensor,
|
|
alpha: float) -> torch.Tensor:
|
|
"""
|
|
Blend two frames with alpha blending.
|
|
|
|
Args:
|
|
frame1: First frame [1, H, W, C]
|
|
frame2: Second frame [1, H, W, C]
|
|
alpha: Blending factor (0.0 = frame1, 1.0 = frame2)
|
|
|
|
Returns:
|
|
Blended frame
|
|
"""
|
|
alpha = torch.clamp(torch.tensor(alpha), 0.0, 1.0)
|
|
return frame1 * (1.0 - alpha) + frame2 * alpha
|
|
|
|
def create_blend_transition(frames1: torch.Tensor, frames2: torch.Tensor,
|
|
blend_length: int) -> torch.Tensor:
|
|
"""
|
|
Create smooth transition between two frame sequences.
|
|
|
|
Args:
|
|
frames1: First sequence [B1, H, W, C]
|
|
frames2: Second sequence [B2, H, W, C]
|
|
blend_length: Number of frames to blend
|
|
|
|
Returns:
|
|
Transition frames [blend_length, H, W, C]
|
|
"""
|
|
if blend_length <= 0:
|
|
return torch.empty(0, frames1.shape[1], frames1.shape[2], frames1.shape[3])
|
|
|
|
# Get last frame of first sequence and first frame of second sequence
|
|
last_frame = frames1[-1:] if frames1.shape[0] > 0 else frames2[:1]
|
|
first_frame = frames2[:1] if frames2.shape[0] > 0 else frames1[-1:]
|
|
|
|
# Ensure frames match dimensions exactly
|
|
if last_frame.shape != first_frame.shape:
|
|
# Resize first_frame to match last_frame exactly
|
|
first_frame = resize_to_match(first_frame, last_frame)
|
|
|
|
# Double check dimensions match
|
|
if last_frame.shape[1:] != first_frame.shape[1:]:
|
|
# Force exact match by creating new tensor with correct dimensions
|
|
target_shape = last_frame.shape
|
|
new_first_frame = torch.zeros_like(last_frame)
|
|
|
|
# Copy data with safe bounds
|
|
min_h = min(first_frame.shape[1], target_shape[1])
|
|
min_w = min(first_frame.shape[2], target_shape[2])
|
|
min_c = min(first_frame.shape[3], target_shape[3])
|
|
|
|
new_first_frame[:, :min_h, :min_w, :min_c] = first_frame[:, :min_h, :min_w, :min_c]
|
|
first_frame = new_first_frame
|
|
|
|
# Create blend sequence
|
|
blend_frames_list = []
|
|
for i in range(blend_length):
|
|
alpha = (i + 1) / (blend_length + 1)
|
|
blended = blend_frames(last_frame, first_frame, alpha)
|
|
blend_frames_list.append(blended)
|
|
|
|
return torch.cat(blend_frames_list, dim=0) if blend_frames_list else torch.empty(0, last_frame.shape[1], last_frame.shape[2], last_frame.shape[3])
|
|
|
|
def safe_frame_index(frame_count: int, index: int, allow_negative: bool = True) -> int:
|
|
"""
|
|
Convert frame index to safe positive index with bounds checking.
|
|
|
|
Args:
|
|
frame_count: Total number of frames
|
|
index: Frame index (can be negative)
|
|
allow_negative: Whether to allow negative indexing
|
|
|
|
Returns:
|
|
Safe positive index
|
|
|
|
Raises:
|
|
ValueError: If index is out of bounds
|
|
"""
|
|
if frame_count <= 0:
|
|
raise ValueError("Frame count must be positive")
|
|
|
|
if allow_negative and index < 0:
|
|
index = frame_count + index
|
|
|
|
if index < 0 or index >= frame_count:
|
|
raise ValueError(f"Frame index {index} out of bounds for {frame_count} frames")
|
|
|
|
return index
|
|
|
|
def chunk_process_frames(frames: torch.Tensor, chunk_size: int = 32) -> list:
|
|
"""
|
|
Split frames into chunks for memory-efficient processing.
|
|
|
|
Args:
|
|
frames: Input frames [B, H, W, C]
|
|
chunk_size: Maximum frames per chunk
|
|
|
|
Returns:
|
|
List of frame chunks
|
|
"""
|
|
if chunk_size <= 0:
|
|
chunk_size = frames.shape[0]
|
|
|
|
chunks = []
|
|
for i in range(0, frames.shape[0], chunk_size):
|
|
chunk = frames[i:i + chunk_size]
|
|
chunks.append(chunk)
|
|
|
|
return chunks
|
|
|
|
|
|
|
|
|