Add fill mode options for frame insertion

- Added new "fill_mode" parameter with options: copy_frame, white, black
- Created helper method to generate solid white or black frames
- Modified processing methods to support different fill modes
- Updated frame insertion logic for both between-frames and first/last frames
- Maintained backward compatibility with original functionality
- Enhanced flexibility for image sequence processing by allowing solid color frames
This commit is contained in:
tok_main01
2025-03-05 17:32:16 +08:00
parent 1666c1db91
commit afddf9756f
2 changed files with 77 additions and 24 deletions
+60 -18
View File
@@ -20,6 +20,7 @@ class QuickImageSequenceProcessNode:
"images": ("IMAGE",), # Input image sequence
"frames_between": ("INT", {"default": 1}), # Number of frames to copy between each frame
"copy_frame": (["previous", "next"],), # Whether to copy the previous or next frame
"fill_mode": (["copy_frame", "white", "black"],), # Fill mode for added frames
"quick_process_first_frames": ("INT", {"default": 0}), # Number of frames to add/remove at the beginning (negative values for removal)
"quick_process_last_frames": ("INT", {"default": 0}), # Number of frames to add/remove at the end (negative values for removal)
"process_F_L_frames_first": (["yes", "no"],), # Whether to prioritize processing first and last frames
@@ -32,51 +33,92 @@ class QuickImageSequenceProcessNode:
CATEGORY = "image/sequence"
def process_between_frames(self, images, frames_between, copy_frame):
def create_solid_frame(self, reference_frame, color="white"):
"""Create a solid white or black frame with the same dimensions as the reference frame"""
# Get dimensions from reference frame
if isinstance(reference_frame, torch.Tensor):
h, w = reference_frame.shape[0], reference_frame.shape[1]
c = reference_frame.shape[2] if len(reference_frame.shape) > 2 else 3
# Create solid frame
if color == "white":
solid_frame = torch.ones((h, w, c), dtype=reference_frame.dtype)
else: # black
solid_frame = torch.zeros((h, w, c), dtype=reference_frame.dtype)
return solid_frame
else:
raise TypeError("Reference frame must be a torch tensor")
def process_between_frames(self, images, frames_between, copy_frame, fill_mode):
"""Process middle frames (including frame insertion)"""
new_images = []
for i in range(len(images)):
# Add current frame
new_images.append(images[i])
# If not the last frame, insert copied frames between current and next
# If not the last frame, insert frames between current and next
if i < len(images) - 1:
if copy_frame == "previous":
frame_to_copy = images[i]
else:
frame_to_copy = images[i + 1]
if fill_mode == "copy_frame":
# Use original copy frame logic
if copy_frame == "previous":
frame_to_add = images[i]
else: # "next"
frame_to_add = images[i + 1]
elif fill_mode == "white":
frame_to_add = self.create_solid_frame(images[i], "white")
else: # "black"
frame_to_add = self.create_solid_frame(images[i], "black")
for _ in range(frames_between):
new_images.append(frame_to_copy)
new_images.append(frame_to_add)
return new_images
def process_end_frames(self, images, quick_process_first_frames, quick_process_last_frames):
def process_end_frames(self, images, quick_process_first_frames, quick_process_last_frames, fill_mode):
"""Process first and last frames (adding or removing)"""
new_images = list(images) # Convert to list for modification
if quick_process_first_frames >= 0:
# Add frames at the beginning
first_frame = new_images[0]
new_images = [first_frame] * quick_process_first_frames + new_images
if fill_mode == "copy_frame":
first_frame = new_images[0]
frames_to_add = [first_frame] * quick_process_first_frames
elif fill_mode == "white":
first_frame = self.create_solid_frame(new_images[0], "white")
frames_to_add = [first_frame] * quick_process_first_frames
else: # "black"
first_frame = self.create_solid_frame(new_images[0], "black")
frames_to_add = [first_frame] * quick_process_first_frames
new_images = frames_to_add + new_images
else:
# Remove frames from the beginning
new_images = new_images[-quick_process_first_frames:]
if quick_process_last_frames >= 0:
# Add frames at the end
last_frame = new_images[-1]
new_images.extend([last_frame] * quick_process_last_frames)
if fill_mode == "copy_frame":
last_frame = new_images[-1]
frames_to_add = [last_frame] * quick_process_last_frames
elif fill_mode == "white":
last_frame = self.create_solid_frame(new_images[-1], "white")
frames_to_add = [last_frame] * quick_process_last_frames
else: # "black"
last_frame = self.create_solid_frame(new_images[-1], "black")
frames_to_add = [last_frame] * quick_process_last_frames
new_images.extend(frames_to_add)
else:
# Remove frames from the end
new_images = new_images[:quick_process_last_frames]
return new_images
def edit_image_sequence(self, images, frames_between, copy_frame, quick_process_first_frames, quick_process_last_frames, process_F_L_frames_first):
def edit_image_sequence(self, images, frames_between, copy_frame, fill_mode, quick_process_first_frames, quick_process_last_frames, process_F_L_frames_first):
"""
Edit image sequence:
1. Process order determined by process_F_L_frames_first
2. Insert specified number of copied frames between frames
2. Insert specified number of frames between frames (copied or solid color)
3. Add or remove frames at the beginning and end
"""
# Ensure input is PyTorch tensor
@@ -85,12 +127,12 @@ class QuickImageSequenceProcessNode:
if process_F_L_frames_first == "yes":
# Process first/last frames first, then middle frames
intermediate_images = self.process_end_frames(images, quick_process_first_frames, quick_process_last_frames)
new_images = self.process_between_frames(intermediate_images, frames_between, copy_frame)
intermediate_images = self.process_end_frames(images, quick_process_first_frames, quick_process_last_frames, fill_mode)
new_images = self.process_between_frames(intermediate_images, frames_between, copy_frame, fill_mode)
else:
# Process middle frames first, then first/last frames
intermediate_images = self.process_between_frames(images, frames_between, copy_frame)
new_images = self.process_end_frames(intermediate_images, quick_process_first_frames, quick_process_last_frames)
intermediate_images = self.process_between_frames(images, frames_between, copy_frame, fill_mode)
new_images = self.process_end_frames(intermediate_images, quick_process_first_frames, quick_process_last_frames, fill_mode)
# Ensure at least one frame remains
if not new_images: