import torch import torch.nn.functional as F class ReplaceFirstLastFrames: @classmethod def INPUT_TYPES(cls): return { "required": { "image_sequence": ("IMAGE",), "num_start_frames": ("INT", { "default": 1, "min": 0, "max": 10000, "step": 1, "display": "number" }), "num_last_frames": ("INT", { "default": 1, "min": 0, "max": 10000, "step": 1, "display": "number" }), }, "optional": { "start_frames": ("IMAGE",), "last_frames": ("IMAGE",), }, } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("image_sequence",) FUNCTION = "replace_frames" CATEGORY = "image/animation" def replace_frames(self, image_sequence, num_start_frames, num_last_frames, start_frames=None, last_frames=None): # Get dimensions from the main video target_height = image_sequence.shape[1] target_width = image_sequence.shape[2] target_channels = image_sequence.shape[3] # Get batch sizes num_images = image_sequence.shape[0] # Edge case: If no frames to replace, return original if num_start_frames <= 0 and num_last_frames <= 0: return (image_sequence,) # Check if start_frames is provided when needed if num_start_frames > 0 and start_frames is None: # If start_frames is not provided but needed, use original frames start_frames = image_sequence[:min(num_start_frames, num_images)] # Check if last_frames is provided when needed if num_last_frames > 0 and last_frames is None: # If last_frames is not provided but needed, use original frames last_frames = image_sequence[max(0, num_images - num_last_frames):] # Resize start_frames and last_frames to match video dimensions if they exist if start_frames is not None: start_frames = self._resize_frames(start_frames, target_height, target_width, target_channels) if last_frames is not None: last_frames = self._resize_frames(last_frames, target_height, target_width, target_channels) # Clamp the number of frames to replace to reasonable values # If the total requested exceeds available frames, adjust proportionally if num_start_frames + num_last_frames > num_images: if num_images == 1: # Special case: only 1 frame total # Prioritize start frame if num_start_frames > 0: start_to_insert = self._prepare_frames(start_frames, 1) return (start_to_insert,) else: last_to_insert = self._prepare_frames(last_frames, 1) return (last_to_insert,) # Adjust to fit available frames total_requested = num_start_frames + num_last_frames ratio = num_images / total_requested # Calculate adjusted values (always keep at least 1 frame in middle if possible) adjusted_start = max(0, int(num_start_frames * ratio)) adjusted_last = max(0, int(num_last_frames * ratio)) # Make sure we don't exceed the total while adjusted_start + adjusted_last >= num_images and num_images > 1: if adjusted_start > adjusted_last: adjusted_start = max(0, adjusted_start - 1) else: adjusted_last = max(0, adjusted_last - 1) # If still too many, just split evenly if adjusted_start + adjusted_last >= num_images: adjusted_start = num_images // 2 adjusted_last = num_images - adjusted_start num_start_frames = adjusted_start num_last_frames = adjusted_last # Prepare start frames to insert start_to_insert = self._prepare_frames(start_frames, num_start_frames) if (num_start_frames > 0 and start_frames is not None) else None # Prepare last frames to insert last_to_insert = self._prepare_frames(last_frames, num_last_frames) if (num_last_frames > 0 and last_frames is not None) else None # Build the output sequence parts = [] # Add start frames if any if start_to_insert is not None: parts.append(start_to_insert) # Add middle frames if any remain middle_start_idx = num_start_frames middle_end_idx = num_images - num_last_frames if middle_end_idx > middle_start_idx: middle_frames = image_sequence[middle_start_idx:middle_end_idx] parts.append(middle_frames) # Add last frames if any if last_to_insert is not None: parts.append(last_to_insert) # Concatenate all parts if len(parts) == 0: # Shouldn't happen, but return original as fallback return (image_sequence,) elif len(parts) == 1: output = parts[0] else: output = torch.cat(parts, dim=0) return (output,) def _resize_frames(self, frames, target_height, target_width, target_channels): # Check if resize is needed if frames.shape[1] == target_height and frames.shape[2] == target_width and frames.shape[3] == target_channels: return frames # ComfyUI format: [batch, height, width, channels] # PyTorch interpolate needs: [batch, channels, height, width] frames_permuted = frames.permute(0, 3, 1, 2) # Resize spatial dimensions if frames.shape[1] != target_height or frames.shape[2] != target_width: frames_permuted = F.interpolate( frames_permuted, size=(target_height, target_width), mode='bilinear', align_corners=False ) # Convert back to ComfyUI format frames_resized = frames_permuted.permute(0, 2, 3, 1) # Handle channel mismatch (e.g., RGB vs RGBA) if frames_resized.shape[3] != target_channels: if target_channels == 3 and frames_resized.shape[3] == 4: # RGBA to RGB: drop alpha frames_resized = frames_resized[:, :, :, :3] elif target_channels == 4 and frames_resized.shape[3] == 3: # RGB to RGBA: add alpha channel (fully opaque) alpha = torch.ones(frames_resized.shape[0], frames_resized.shape[1], frames_resized.shape[2], 1, dtype=frames_resized.dtype, device=frames_resized.device) frames_resized = torch.cat([frames_resized, alpha], dim=3) elif frames_resized.shape[3] == 1 and target_channels == 3: # Grayscale to RGB: repeat channel frames_resized = frames_resized.repeat(1, 1, 1, 3) elif frames_resized.shape[3] == 1 and target_channels == 4: # Grayscale to RGBA: repeat channel + add alpha frames_resized = frames_resized.repeat(1, 1, 1, 3) alpha = torch.ones(frames_resized.shape[0], frames_resized.shape[1], frames_resized.shape[2], 1, dtype=frames_resized.dtype, device=frames_resized.device) frames_resized = torch.cat([frames_resized, alpha], dim=3) return frames_resized def _prepare_frames(self, source_frames, num_needed): if num_needed <= 0: return None num_available = source_frames.shape[0] if num_available == num_needed: # Perfect match return source_frames elif num_available == 1: # Repeat single frame return source_frames.repeat(num_needed, 1, 1, 1) elif num_available > num_needed: # More frames than needed, use the first ones return source_frames[:num_needed] else: # Fewer frames than needed, cycle through them # Calculate how many full cycles + remaining frames full_cycles = num_needed // num_available remainder = num_needed % num_available if remainder == 0: # Perfect multiple return source_frames.repeat(full_cycles, 1, 1, 1) else: # Need partial cycle full_part = source_frames.repeat(full_cycles, 1, 1, 1) partial_part = source_frames[:remainder] return torch.cat([full_part, partial_part], dim=0) # ComfyUI node registration NODE_CLASS_MAPPINGS = { "ReplaceFirstLastFrames": ReplaceFirstLastFrames } NODE_DISPLAY_NAME_MAPPINGS = { "ReplaceFirstLastFrames": "Replace First & Last Frames" }