Files
TashaSkyUp-EternalKernelPyt…/AVF/video_nodes.py
T
2024-11-11 18:27:04 -07:00

696 lines
26 KiB
Python

# video_module.py
from .common import ETK_AVF_base, duplicate_file, run_ffmpeg_command
@ETK_AVF_base
class VideoFolderFFMPEGVideoInterpolate:
"""
Use run_ffmpeg to interpolate the frames in the video folder
"""
@classmethod
def INPUT_TYPES(cls):
ret = {
"required": {
"video_folder_def": ("VIDEO_FOLDER_DEF", {"default": None}),
"last_n_frames": ("INT", {"default": 2}),
"target_n_frames": ("INT", {"default": 3}),
},
}
return ret
RETURN_TYPES = ("VIDEO_FOLDER_DEF",)
RETURN_NAMES = ("video_folder_def",)
FUNCTION = "video_folder_ffmpeg_video_interpolate"
def video_folder_ffmpeg_video_interpolate(self, **kwargs):
import os
import shutil
video_folder_def = kwargs["video_folder_def"]
last_n_frames = kwargs["last_n_frames"]
target_n_frames = kwargs["target_n_frames"]
# Get the current frame count
current_frame_count = video_folder_def.get_frame_count()
# Calculate the target frame rate for interpolation
target_fps = video_folder_def["fps"] * (target_n_frames / last_n_frames)
source_fps = video_folder_def["fps"]
# Get the file names for the last n frames
last_frame_names = [video_folder_def.get_last_file_name(offset=i) for i in range(last_n_frames)]
last_frame_names = list(set(last_frame_names))
last_frame_names.sort()
# Create a temporary folder off of the video frame folder's base path
temp_folder = os.path.join(video_folder_def["video_folder"], "temp")
os.makedirs(temp_folder, exist_ok=True)
# Move the input frames into the new folder
for frame_name in last_frame_names:
# Do this in a way that overwrites the files and MOVES them
shutil.move(frame_name, os.path.join(temp_folder, os.path.basename(frame_name)))
# New input frame names list
last_frame_names = [os.path.join(temp_folder, os.path.basename(frame_name)) for frame_name in last_frame_names]
last_frame_names.sort()
# Create a file with the list of input frame paths in the temporary folder
input_frame_list_file = os.path.join(temp_folder, "input_frames.txt")
input_frame_list_file = os.path.abspath(input_frame_list_file)
with open(input_frame_list_file, 'w') as file:
# Need to do this in a manner that puts the first frame at the top of the file.
file.write('\n'.join([f"file '{frame_name}'" for frame_name in last_frame_names]))
# Get the output frame names for the interpolated frames
output_frame_names = [video_folder_def.get_next_file_name(i) for i in range(target_n_frames)]
temp_output_frame_pattern = os.path.join(video_folder_def["video_folder"], "temp",
"tmp_frame_%010d." + video_folder_def["frame format"].lower())
temp_output_frame_pattern = os.path.abspath(temp_output_frame_pattern)
# Create the FFmpeg command for motion interpolation using the frame list file
ffmpeg_interpolate_cmd = [
"-f", "concat", "-safe", "0",
"-r", str(source_fps),
"-i", input_frame_list_file,
"-vsync", "1",
"-filter_complex",
f"[0:v]setpts=N/TB,minterpolate=fps={target_fps}:mi_mode=blend,setpts=PTS-STARTPTS[v]", # Reset timestamps
"-map", "[v]", "-q:v", "2",
"-r", str(target_fps),
"-frames:v", str(target_n_frames),
"-threads", "0",
"-y", temp_output_frame_pattern
]
run_ffmpeg_command(ffmpeg_interpolate_cmd)
# Move the interpolated frames to the video folder with the correct names
for i, frame_name in enumerate(output_frame_names):
temp_frame_name = os.path.join(temp_folder,
f"tmp_frame_{i + 1:010d}.{video_folder_def['frame format'].lower()}")
try:
shutil.move(temp_frame_name, frame_name)
except FileNotFoundError:
pass
# Remove the temporary folder and frame list file
os.remove(input_frame_list_file)
# Remove this without generating an error even if it is not empty
try:
os.rmdir(temp_folder)
except OSError:
shutil.rmtree(temp_folder)
return (video_folder_def,)
@ETK_AVF_base
class VideoDefinitionProvider:
"""
Given a VIDEOFOLDER node, and inputs for height, width, and fps, creates a VIDEOFOLDERDEF variable
"""
@classmethod
def INPUT_TYPES(cls):
ret = {
"required": {
"video_folder": ("VIDEOFOLDER", {"default": None}),
"height": ("INT", {"default": 1080}),
"width": ("INT", {"default": 1920}),
"fps": ("FLOAT", {"default": 30}),
"frame format": (["JPG", "PNG"], {"default": "JPG"}),
},
}
return ret
RETURN_TYPES = ("VIDEO_FOLDER_DEF",)
FUNCTION = "video_definition_provider"
def video_definition_provider(self, **kwargs):
from addict import Dict as adict
import os
ret = adict()
ret["video_folder"] = kwargs["video_folder"]
ret["height"] = kwargs["height"]
ret["width"] = kwargs["width"]
ret["fps"] = kwargs["fps"]
ret["frame format"] = kwargs.get("frame format") or kwargs.get("frame_format")
# A function that gets the frame count from the disk
ret.get_frame_count = lambda: len([f for f in os.listdir(ret["video_folder"]) if "frame_" in f])
# A function that gets the next file name and full path, don't forget to add 10 zeros to the frame count
ret.get_next_file_name = lambda x=0: os.path.join(
ret['video_folder'],
f"frame_{str(ret.get_frame_count() + x).zfill(10)}.{ret['frame format'].lower()}"
)
ret.get_last_file_name = lambda offset=0: os.path.join(
ret['video_folder'],
f"frame_{str(max(ret.get_frame_count() - offset - 1, 0)).zfill(10)}.{ret['frame format'].lower()}"
)
ret.get_current_file_name = lambda: os.path.join(
ret['video_folder'],
f"frame_{str(ret.get_frame_count() - 1).zfill(10)}.{ret['frame format'].lower()}"
)
def reserve_next_frame():
fn = ret.get_next_file_name()
open(fn, "w").close()
return fn
ret.reserve_next_frame = reserve_next_frame
return (ret,)
@ETK_AVF_base
class VideoFolderProvider:
"""
Given a folder path, creates the folder or empties it entirely returns the folder path as
a string but as the node type "VIDEOFOLDER"
"""
@classmethod
def INPUT_TYPES(cls):
ret = {
"required": {
"folder_path": ("STRING", {"default": ""}),
"clear_folder": ("BOOLEAN", {"default": True}),
},
"optional": {
"activator:": ("*", {"default": None})
}
}
return ret
RETURN_TYPES = ("VIDEOFOLDER",)
FUNCTION = "video_folder_provider"
def video_folder_provider(self, **kwargs):
import shutil
import os
# Check if the folder_path is none
if kwargs["folder_path"] is None:
raise ValueError("folder_path must be specified")
# Set folder exists flag by using the os.path.exists function
folder_exists = os.path.exists(kwargs["folder_path"])
if not folder_exists:
os.makedirs(kwargs["folder_path"])
folder_exists = True
# Set files in folder flag by using the os.listdir function
files_in_folder = os.listdir(kwargs["folder_path"])
# Set accessible flag by using the os.access function
accessible = os.access(kwargs["folder_path"], os.W_OK)
if not accessible:
raise ValueError("Folder for video files is not accessible")
if kwargs["clear_folder"]:
# If the folder exists, delete all files in the folder
if folder_exists:
for file in files_in_folder:
# Use shutil to remove the file
fl = os.path.join(kwargs["folder_path"], file)
if os.path.isfile(fl):
os.remove(fl)
# Use shutil to remove the entire path
check_this_path = file
if len(check_this_path) > 3:
shutil.rmtree(check_this_path, ignore_errors=True)
else:
raise ValueError("Folder path is too short")
return (kwargs["folder_path"],)
@ETK_AVF_base
class AddFramesToVideoFolder:
"""
Given a VIDEO_FOLDER_DEF node, and a list of frames, adds the frames to the video folder
"""
@classmethod
def INPUT_TYPES(cls):
ret = {
"required": {
"video_folder_def": ("VIDEO_FOLDER_DEF", {"default": None}),
"IMAGE": ("IMAGE", {"default": None}),
"Rescale method": (
["NEAREST", "BOX", "BILINEAR", "HAMMING", "BICUBIC", "LANCZOS"],
{"default": "LANCZOS"}),
"Release vram?": ("BOOLEAN", {"default": True}),
},
}
return ret
RETURN_TYPES = ("VIDEO_FOLDER_DEF",)
FUNCTION = "add_frames_to_video_folder"
@staticmethod
def save_or_copy_frame(current_frame, previous_frame, frame_fp, previous_frame_fp):
import shutil
from PIL import Image
import torch
if previous_frame is None or not torch.equal(current_frame, previous_frame):
# If the current frame is different from the previous frame, save it to disk
current_frame = current_frame
# expected shape is H, W, C uint8 0-255
# image.fromarray expects uint8 0-255 and shape H, W, C
needs_dtype_conversion = current_frame.dtype != torch.uint8
if len(current_frame.shape) == 3 and current_frame.dtype == torch.uint8:
needs_dtype_conversion = False
elif len(current_frame.shape) == 2:
# assume grayscale H,W
current_frame = current_frame.unsqueeze(2).repeat(1, 1, 3)
elif len(current_frame.shape) == 4 and current_frame.shape[0] == 1:
# assume B, H, W, C
current_frame = current_frame.squeeze(0)
elif len(current_frame.shape) == 3 and current_frame.shape[-1] == 4:
# assume H, W, C with alpha
current_frame = current_frame[..., :3]
else: # try to move forward
pass
# assume H, W, C
if needs_dtype_conversion:
current_frame = current_frame.mul(255).byte()
ndarr = current_frame.cpu().numpy()
im = Image.fromarray(ndarr)
im.save(frame_fp)
previous_frame = current_frame
else:
# If the current frame is the same as the previous frame, copy the previously saved frame file
shutil.copy2(previous_frame_fp, frame_fp)
return previous_frame
@staticmethod
def process_frames(frames, rescale_method, video_folder_def):
import gc
import torch
num_frames = frames.shape[0]
previous_frame = None
frame_fp = None
frames_are_all_same = torch.all(frames[0] == frames)
if frames_are_all_same and frames.shape[0] > 1:
# Get a list of file names using reserve_next_frame
frame_f_names = [video_folder_def.reserve_next_frame() for _ in range(num_frames)]
# Save the first frame
AddFramesToVideoFolder.save_or_copy_frame(frames[0], previous_frame, frame_f_names[0], None)
# Copy the first frame to the rest of the frames using os by calling the cli
duplicate_file(frame_f_names[0], frame_f_names[1:])
else:
for i in range(num_frames):
next_fn = video_folder_def.get_next_file_name()
previous_frame_fp = frame_fp
frame_fp = next_fn
current_frame = frames[i]
previous_frame = AddFramesToVideoFolder.save_or_copy_frame(current_frame, previous_frame, frame_fp,
previous_frame_fp)
del frames
gc.collect()
return None
def add_frames_to_video_folder(self, **kwargs):
from PIL import Image
import torch
import gc
from torchvision.transforms.functional import resize
import psutil
from custom_nodes.EternalKernelLiteGraphNodes.modules.image_utils import lanczos_resize
initial_memory = psutil.virtual_memory().available
print(f"Initial free memory: {initial_memory / (1024 * 1024):.2f} MB")
# Get the video folder path
video_folder = kwargs["video_folder_def"]["video_folder"]
# Get the frame format
frame_format = kwargs["video_folder_def"]["frame format"]
# Get the height
height = kwargs["video_folder_def"]["height"]
# Get the width
width = kwargs["video_folder_def"]["width"]
# Get the fps
fps = kwargs["video_folder_def"]["fps"]
# Get the frames
frames = kwargs["IMAGE"]
# Get the rescale method
rescale_method = kwargs.get("Rescale method", None)
if not rescale_method:
rescale_method = kwargs.get("Rescale_method", None)
free_vram = kwargs.get("Release vram?", None)
if not free_vram:
free_vram = kwargs.get("Release_vram_", True)
orig_frames = frames
# num_files = len(glob(os.path.join(video_folder, f"frame_*.{frame_format.lower()}")))
# Rescale if needed to the correct size
if height != frames.shape[1] or width != frames.shape[2]:
# First since transforms.Resize requires a tensor of shape B,C, H, W
# We need to permute the tensor to B, C, H, W
frames = frames.permute(0, 3, 1, 2)
frames = frames.to("cuda")
# Get the right value for interpolation
interpolation_methods = {
"NEAREST": Image.Resampling.NEAREST,
"BOX": Image.Resampling.BOX,
"BILINEAR": Image.Resampling.BILINEAR,
"HAMMING": Image.Resampling.HAMMING,
"BICUBIC": Image.Resampling.BICUBIC,
"LANCZOS": Image.Resampling.LANCZOS
}
rescale_type = interpolation_methods[rescale_method]
# rescale_method = transforms.Resize((height, width), interpolation=rescale_type) # this doesnt work so well lets use a different function
if rescale_type != Image.Resampling.LANCZOS:
rescale_method_fn = lambda x: resize(x, (height, width), interpolation=rescale_type)
frames = rescale_method_fn(frames)
else:
rescale_method_fn = lambda x: lanczos_resize(x, (width, height))
frames = rescale_method_fn(frames)
# Now we need to permute back to B, H, W, C
frames = frames.permute(0, 2, 3, 1)
self.process_frames(frames, rescale_method_fn, kwargs["video_folder_def"])
else:
self.process_frames(frames, rescale_method, kwargs["video_folder_def"])
gc.collect()
if free_vram:
import gc
if isinstance(frames, torch.Tensor):
frames.to("cpu").detach().numpy()
del frames
if isinstance(orig_frames, torch.Tensor):
orig_frames = orig_frames.to("cpu").detach().numpy()
del orig_frames
torch.cuda.empty_cache()
gc.collect()
initial_memory = psutil.virtual_memory().available
print(f"After adding frames free memory: {initial_memory / (1024 * 1024):.2f} MB")
# Return the VIDEO_FOLDER_DEF
return (kwargs["video_folder_def"],)
@ETK_AVF_base
class HoldFrameForSecsInVideoFolder:
"""
Given a VIDEO_FOLDER_DEF node, and a list of frames, first uses RepeatFrames to repeat the frames
for the correct amount and then adds the frames to the video folder using AddFramesToVideoFolder
"""
@classmethod
def INPUT_TYPES(cls):
ret = {
"required": {
"video_folder_def": ("VIDEO_FOLDER_DEF", {"default": None}),
"IMAGE": ("IMAGE", {"default": None}),
"seconds": ("FLOAT", {"default": 1.0}),
"Rescale method": (
["NEAREST", "BOX", "BILINEAR", "HAMMING", "BICUBIC", "LANCZOS"],
{"default": "LANCZOS"}),
"Release vram?": ("BOOLEAN", {"default": True}),
},
}
return ret
RETURN_TYPES = ("VIDEO_FOLDER_DEF",)
FUNCTION = "hold_frame_for_secs_in_video_folder"
def saver(self, x, kwargs):
# This just saves in an async way
rescale_method = kwargs.get("Rescale method", None)
if rescale_method is None:
rescale_method = kwargs.get("Rescale_method", None)
rel_vram = kwargs.get("Release vram?", None)
if rel_vram is None:
rel_vram = kwargs.get("Release_vram_", True)
AddFramesToVideoFolder().add_frames_to_video_folder(
video_folder_def=kwargs["video_folder_def"],
IMAGE=x,
Rescale_method=rescale_method,
Release_vram_=rel_vram
)
def hold_frame_for_secs_in_video_folder(self, **kwargs):
video_folder_ref = kwargs["video_folder_def"]
video_folder = kwargs["video_folder_def"]["video_folder"]
# Get the frame format
frame_format = kwargs["video_folder_def"]["frame format"]
# Get the height
height = kwargs["video_folder_def"]["height"]
# Get the width
width = kwargs["video_folder_def"]["width"]
# Get the fps
fps = kwargs["video_folder_def"]["fps"]
# Get the frames
frames = kwargs["IMAGE"]
# Get the number of frames
num_frames = frames.shape[0]
# Get the rescale method
rescale_method = kwargs.get("Rescale method", None)
if rescale_method is None:
rescale_method = kwargs.get("Rescale_method", None)
# Get the number of seconds
seconds = kwargs["seconds"]
# Get the number of frames to hold
num_frames_to_hold = int(fps * seconds)
num_frames_in_this_stack = frames.shape[0]
# Get the number of repeats needed
num_repeats_needed = num_frames_to_hold // num_frames_in_this_stack
# Get the number of frames left over
num_frames_left_over = num_frames_to_hold % num_frames_in_this_stack
# Write the first set of frames to disk
self.saver(frames, kwargs)
# Get the name of the last frame
source_last_frame_name = video_folder_ref.get_last_file_name()
# From the last frame find the rest of the frames names/paths
source_last_frame_num = int(source_last_frame_name.split("_")[-1].split(".")[0])
source_first_frame_num = source_last_frame_num - num_frames_in_this_stack + 1
template = source_last_frame_name.replace(str(source_last_frame_num).zfill(10), "{}")
range_of_source_frame_numbers = list(range(source_first_frame_num, source_last_frame_num + 1))
source_frame_names = [template.format(str(x).zfill(10)) for x in range_of_source_frame_numbers]
target_done_frame_num = source_last_frame_num + num_frames_to_hold - num_frames_in_this_stack
num_frames_remaining = num_frames_to_hold - num_frames_in_this_stack
target_start_frame_num = source_last_frame_num + 1
target_frame_range = list(range(target_start_frame_num, target_done_frame_num + 1))
target_frame_names = [template.format(str(x).zfill(10)) for x in target_frame_range]
set_todo = {fn: [] for fn in source_frame_names}
# We can now repeat the frames
for i, target_frame_name in enumerate(target_frame_names):
# Keep track of what stack frame we are on
stack_frame_idx = i % num_frames_in_this_stack
stack_frame_name_to_use = source_frame_names[stack_frame_idx]
set_todo[stack_frame_name_to_use].append(target_frame_name)
# Now we can repeat the frames
for frame_name in set_todo:
src_frame = frame_name
target_frames = set_todo[src_frame]
duplicate_file(src_frame, target_frames)
return (kwargs["video_folder_def"],)
@ETK_AVF_base
class Listed_HoldFramesForSecsInVideoFolder:
"""This just calls HoldFrameForSecsInVideoFolder multiple times for a list of frames and a list of durations"""
@classmethod
def INPUT_TYPES(cls):
ret = {
"required": {
"video_folder_def": ("VIDEO_FOLDER_DEF", {"default": None}),
"IMAGE": ("LIST", {"default": None}),
"seconds": ("LIST", {"default": None}),
"Rescale method": (
["NEAREST", "BOX", "BILINEAR", "HAMMING", "BICUBIC", "LANCZOS"],
{"default": "LANCZOS"}),
"Release vram?": ("BOOLEAN", {"default": True}),
},
}
return ret
RETURN_TYPES = ("VIDEO_FOLDER_DEF",)
FUNCTION = "listed_hold_frames_for_secs_in_video_folder"
def listed_hold_frames_for_secs_in_video_folder(self, **kwargs):
# The plan is to call HoldFrameForSecsInVideoFolder for each frame and duration
# and then return the video_folder_def
if "IMAGE" not in kwargs or "seconds" not in kwargs:
return (kwargs["video_folder_def"],)
for frame, duration in zip(kwargs["IMAGE"], kwargs["seconds"]):
HoldFrameForSecsInVideoFolder().hold_frame_for_secs_in_video_folder(
video_folder_def=kwargs["video_folder_def"],
IMAGE=frame,
seconds=duration,
Rescale_method=kwargs["Rescale method"],
Release_vram_=kwargs["Release vram?"])
return (kwargs["video_folder_def"],)
@ETK_AVF_base
class VideoFolderDefInfo:
"""
Returns everything in the VIDEO_FOLDER_DEF, and also returns information about the status of the folder
"""
@classmethod
def INPUT_TYPES(cls):
ret = {
"required": {
"video_folder_def": ("VIDEO_FOLDER_DEF", {"default": None}),
},
}
return ret
RETURN_TYPES = ("VIDEO_FOLDER_DEF", "STRING", "INT", "INT", "FLOAT", "STRING", "INT", "LIST",)
RETURN_NAMES = ("VFD", "video_folder", "height", "width", "fps", "frame format", "num_files", "list of files")
OUTPUT_NODE = True
FUNCTION = "video_folder_def_info"
def video_folder_def_info(self, **kwargs):
import os
video_folder = kwargs["video_folder_def"]["video_folder"]
height = kwargs["video_folder_def"]["height"]
width = kwargs["video_folder_def"]["width"]
fps = kwargs["video_folder_def"]["fps"]
frame_format = kwargs["video_folder_def"]["frame format"]
files = os.listdir(video_folder)
num_files = len(files)
return (kwargs["video_folder_def"], video_folder, height, width, fps, frame_format, num_files, files,)
@ETK_AVF_base
class FadeToInVideoFolder:
"""
Given a VIDEO_FOLDER_DEF node, and a stack of frames, fades the frames in and adds them to the video folder
"""
@classmethod
def INPUT_TYPES(cls):
ret = {
"required": {
"video_folder_def": ("VIDEO_FOLDER_DEF", {"default": None}),
"IMAGE": ("IMAGE", {"default": None}),
"seconds": ("FLOAT", {"default": 1.0}),
"transition ms": ("FLOAT", {"default": 0.1}),
"Rescale method": (
["NEAREST", "BOX", "BILINEAR", "HAMMING", "BICUBIC", "LANCZOS"],
{"default": "LANCZOS"}),
"Release vram?": ("BOOLEAN", {"default": True}),
},
}
return ret
RETURN_TYPES = ("VIDEO_FOLDER_DEF",)
FUNCTION = "fade_to_in_video_folder"
def fade_to_in_video_folder(self, **kwargs):
from PIL import Image
import torch
# Use HoldFrameForSecsInVideoFolder to hold the frame after creating and AddFramesToVideoFolder the transition
transition_length_secs = int(kwargs["transition ms"] * kwargs["video_folder_def"]["fps"]) / 1000
hold_frame_length_secs = int(kwargs["seconds"] * kwargs["video_folder_def"]["fps"]) - transition_length_secs
transition_frame_count = int(kwargs["video_folder_def"]["fps"] * transition_length_secs)
base_frame_fn = kwargs["video_folder_def"].get_last_file_name()
# Load it into a pytorch tensor (B, C, H, W)
base_frame = Image.open(base_frame_fn)
base_frame = torch.tensor(base_frame).permute(2, 0, 1).unsqueeze(0).to("cuda")
transition_stack = torch.zeros(
(transition_frame_count, base_frame.shape[1], base_frame.shape[2], base_frame.shape[3]),
dtype=torch.float16, device="cuda"
)
# Fade from base_frame to the first frame in the stack over the transition length
new_frame = kwargs["IMAGE"][0].to("cuda")
for i in range(transition_frame_count):
alpha = i / transition_frame_count
interpolated_frame = (1 - alpha) * base_frame + alpha * new_frame
interpolated_frame = interpolated_frame.to("cpu")
transition_stack[i] = interpolated_frame
# Add to the video folder
AddFramesToVideoFolder().add_frames_to_video_folder(
video_folder_def=kwargs["video_folder_def"],
IMAGE=transition_stack,
Rescale_method=kwargs["Rescale method"],
Release_vram_=kwargs["Release vram?"]
)
# Hold the frame
HoldFrameForSecsInVideoFolder().hold_frame_for_secs_in_video_folder(
video_folder_def=kwargs["video_folder_def"],
IMAGE=kwargs["IMAGE"],
seconds=hold_frame_length_secs,
Rescale_method=kwargs["Rescale method"],
Release_vram_=kwargs["Release vram?"]
)
return (kwargs["video_folder_def"],)