Tasha Upchurch
d9170f804b
Key updates: Added AddFramesToVideoFolder node to add frames to a video folder Introduced VideoFolderProvider and VideoDefinitionProvider nodes to manage video folders and video definitions Added FadeToInVideoFolder node for fading between frames in a video folder Implemented ListedSVDConditioning node for applying SVD conditioning to a list of parameters Added VideoFolderFFMPEGVideoInterpolate node to interpolate frames in a video folder using FFmpeg Introduced AudioFolderProvider and AudioDefinitionProvider nodes for managing audio folders and definitions Added AddAudioToAudioFolder and AddSilenceToAudioFolder nodes for adding audio files and silence to audio folders Implemented GetAudioFilesLengthsFromAudioFolder node to retrieve lengths of audio files in a folder Added AudioFolderRenderAudio node to join audio files in an audio folder Introduced GenerateMusic node for generating music from descriptions using the AudioCraft library Added XttsNode for text-to-speech conversion using the XTTS library Implemented AdjustAudioVolumeForFiles node to adjust volume for a list of audio files The changes also include the addition of LiteLLM nodes for text generation: Added FuncAnysToText node that runs Python code with multiple input types and outputs a string Introduced FuncAnysToLatent node for running code that outputs a latent representation Added FuncAnysToString node for executing code with various inputs and producing a string output Implemented FuncAnysToList node to run code that takes different input types and returns a list Added FuncAnysToCond node for running code with multiple inputs and generating a conditioning output These new nodes expand the capabilities for video and audio processing, allowing more advanced manipulations and generation of video and audio content within the graph.
3129 lines
116 KiB
Python
3129 lines
116 KiB
Python
### info for code completion AI ###
|
|
"""
|
|
all of these classes are plugins for comfyui and follow the same pattern
|
|
all of the images are torch tensors and it is unknown and unimportant if they are on the cpu or gpu
|
|
all image inputs are (B,W,H,C)
|
|
|
|
avoid numpy and PIL as much as possible
|
|
"""
|
|
import shutil
|
|
|
|
import fractions
|
|
import os
|
|
from abc import ABC, ABCMeta, abstractmethod
|
|
from typing import List, Tuple, Union
|
|
import cv2
|
|
import torch
|
|
import torchaudio
|
|
# from cv2 import VideoWriter, VideoWriter_fourcc, VideoCapture
|
|
# from line_profiler import profile
|
|
from torch import dtype
|
|
#from custom_nodes.EternalKernelLiteGraphNodes.image import torch_image_show
|
|
from torchvision.models.optical_flow import raft_large, Raft_Large_Weights, raft_small, Raft_Small_Weights
|
|
from custom_nodes.EternalKernelLiteGraphNodes.modules.image_utils import lanczos_resize
|
|
import os
|
|
import torch
|
|
from torchvision.utils import draw_keypoints
|
|
from skimage.metrics import structural_similarity as ssim
|
|
import cv2
|
|
import numpy as np
|
|
import cv2
|
|
import subprocess
|
|
import imageio_ffmpeg as ffmpeg
|
|
from comfy_extras.nodes_video_model import SVD_img2vid_Conditioning
|
|
|
|
NODE_CLASS_MAPPINGS = {} # this is the dictionary that will be used to register the nodes
|
|
|
|
# setup directories
|
|
root_p = os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
|
|
# input_dir = os.path.join(root_p, 'input', 'video')
|
|
# output_dir = os.path.join(root_p, 'output', 'video')
|
|
|
|
video_dir = os.path.join(root_p, 'video')
|
|
|
|
if not os.path.exists(video_dir):
|
|
os.makedirs(video_dir)
|
|
|
|
|
|
# video_folders = lambda s: sorted(
|
|
# [name for name in os.listdir(s.input_dir) if os.path.isdir(os.path.join(s.input_dir, name))])
|
|
# turn the above in the a full function
|
|
def video_folders(s):
|
|
ret = sorted(
|
|
[name for name in os.listdir(s.input_dir) if os.path.isdir(os.path.join(s.input_dir, name))])
|
|
return ret
|
|
|
|
|
|
def video_files(s):
|
|
# any files in in_video_folders ending in .mp4
|
|
ret = sorted([name for name in os.listdir(s.input_dir) if
|
|
os.path.isfile(os.path.join(s.input_dir, name)) and name.endswith('.mp4')])
|
|
return ret
|
|
|
|
|
|
def run_ffmpeg_command(args):
|
|
"""
|
|
Executes a command using the ffmpeg executable.
|
|
|
|
Parameters:
|
|
args (list): List of command-line arguments for ffmpeg.
|
|
"""
|
|
if isinstance(args, str):
|
|
args = args.split()
|
|
cmd = [ffmpeg.get_ffmpeg_exe()] + args
|
|
subprocess.run(cmd)
|
|
|
|
|
|
# start with an ABC to define the common widget interface
|
|
|
|
class ABCWidgetMetaclass(ABCMeta):
|
|
"""A metaclass that automatically registers classes."""
|
|
|
|
def __init__(cls, name, bases, attrs):
|
|
if (ABC in bases) or (ABCVideoWidget in bases) or ("ABC" in name):
|
|
pass
|
|
else:
|
|
NODE_CLASS_MAPPINGS[str(name)] = cls
|
|
cls.CATEGORY = "ETK/video"
|
|
super().__init__(name, bases, attrs)
|
|
|
|
|
|
class ABCVideoWidget(ABC, metaclass=ABCWidgetMetaclass):
|
|
"""Abstract base class for simple construction"""
|
|
|
|
input_dir = video_dir
|
|
output_dir = video_dir
|
|
|
|
CATEGORY = "video"
|
|
RETURN_TYPES = ("VIDEO",)
|
|
FUNCTION = "handler"
|
|
|
|
@classmethod
|
|
@abstractmethod
|
|
def INPUT_TYPES(cls):
|
|
"""All subclasses must provide a INPUT_TYPES classmethod"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def handler(self, one: object, two: object):
|
|
"""All subclasses must provide a handler method"""
|
|
pass
|
|
|
|
|
|
class ABCVideoFileToABCVideoFolder(ABCVideoWidget, metaclass=ABCWidgetMetaclass):
|
|
"""Abstract base class for simple construction"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required":
|
|
{
|
|
"video_in": (video_files(cls),),
|
|
"folder_out": (video_folders(cls),),
|
|
},
|
|
|
|
"optional":
|
|
{"text":
|
|
("STRING", {"multiline": False}),
|
|
}
|
|
}
|
|
|
|
@abstractmethod
|
|
def handler(self, one, two, text):
|
|
return (one, two,)
|
|
|
|
|
|
class ABCABCVideoFolderToImage(ABCVideoWidget, metaclass=ABCWidgetMetaclass):
|
|
"""Abstract base class for simple construction"""
|
|
RETURN_TYPES = ("IMAGE",)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required":
|
|
{
|
|
"folder_in": ("STRING", {"default": "full_path_to_folder"}),
|
|
"idx_slice_start": ("INT", {"min": 0, "max": 99999, "step": 1, "default": 0}),
|
|
"idx_slice_stop": ("INT", {"min": 0, "max": 99999, "step": 1, "default": 1}),
|
|
"slice_idx_step": ("INT", {"min": 0, "max": 9999, "step": 1, "default": 1}),
|
|
"idx": ("INT", {"min": 0, "max": 99999, "step": 1, "default": 0}),
|
|
},
|
|
|
|
"optional":
|
|
{"text":
|
|
("STRING", {"multiline": False}),
|
|
}
|
|
}
|
|
|
|
@abstractmethod
|
|
def handler(self, one, two, text):
|
|
return (one,)
|
|
|
|
|
|
class ABCABCVideoFileToImage(ABCVideoWidget, metaclass=ABCWidgetMetaclass):
|
|
"""Abstract base class for simple construction"""
|
|
RETURN_TYPES = ("IMAGE",)
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required":
|
|
{
|
|
"video_in": ("STRING", {"default": "full_path_to_video"}),
|
|
"idx_slice_start": ("INT", {"min": 0, "max": 99999, "step": 1, "default": 0}),
|
|
"idx_slice_stop": ("INT", {"min": 0, "max": 99999, "step": 1, "default": 1}),
|
|
"slice_idx_step": ("INT", {"min": 0, "max": 9999, "step": 1, "default": 1}),
|
|
"idx": ("INT", {"min": 0, "max": 99999, "step": 1, "default": 0}),
|
|
"func_only": ([False, True], {"default": False}),
|
|
},
|
|
|
|
}
|
|
|
|
@abstractmethod
|
|
def handler(self, one, two, text):
|
|
return (one,)
|
|
|
|
|
|
## end of abstract base classes ##
|
|
|
|
class VideoToFramesFolder(ABCVideoFileToABCVideoFolder, metaclass=ABCWidgetMetaclass):
|
|
"""use cv2 to open the video and use it as a source for image stacks"""
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.old = self.INPUT_TYPES
|
|
# self.INPUT_TYPES = lambda s: self.old().update({})
|
|
|
|
def handler(self, video_in, folder_out, text):
|
|
import time
|
|
start = time.time()
|
|
video_in_path = os.path.join(self.input_dir, video_in)
|
|
video_out_path = os.path.join(self.output_dir, folder_out)
|
|
|
|
vidcap = cv2.VideoCapture(video_in_path)
|
|
success, image = vidcap.read()
|
|
count = 0
|
|
while success:
|
|
cv2.imwrite(os.path.join(video_out_path, "frame%d.jpg" % count), image)
|
|
success, image = vidcap.read()
|
|
count += 1
|
|
|
|
vidcap.release()
|
|
end = time.time()
|
|
print("time elapsed: ", end - start)
|
|
return (folder_out,)
|
|
|
|
|
|
class VideoToFramesFolderFFMPEG(ABCVideoFileToABCVideoFolder, metaclass=ABCWidgetMetaclass):
|
|
"""use ffmpeg to open the video and use it as a source for image stacks, use PNG or jpeg with quality 100"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required":
|
|
{
|
|
"video_in": ("STRING", {"default": "full_path_to_folder"}),
|
|
"folder_out": ("STRING", {"default": "full_path_to_folder"}),
|
|
"format": (["PNG", "JPEG"], {"default": "PNG"}),
|
|
},
|
|
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("folder_out",)
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.old = self.INPUT_TYPES
|
|
# self.INPUT_TYPES = lambda s: self.old().update({})
|
|
|
|
def handler(self, video_in, folder_out, format):
|
|
import time
|
|
start = time.time()
|
|
video_in_path = os.path.join(self.input_dir, video_in)
|
|
video_out_path = os.path.join(self.output_dir, folder_out)
|
|
|
|
if not os.path.exists(video_out_path):
|
|
os.makedirs(video_out_path)
|
|
|
|
# specifically this calls ffmpeg with the following arguments:
|
|
# -i means input file
|
|
# -qscale:v means quality scale video, default is 2, which is a good compromise between size and quality
|
|
# 1 is the quality scale value, which is a value between 1 and 31, where 1 is the highest quality
|
|
|
|
start = time.time()
|
|
|
|
if format == "JPEG":
|
|
subprocess.call(
|
|
["ffmpeg", "-i", video_in_path, "-qscale:v", "1", os.path.join(video_out_path, "frame%05d.jpg")])
|
|
end = time.time()
|
|
print("ffmpeg jpg process took {} seconds".format(end - start))
|
|
|
|
elif format == "PNG":
|
|
subprocess.call(["ffmpeg", "-i", video_in_path, os.path.join(video_out_path, "frame%05d.png")])
|
|
end = time.time()
|
|
print("ffmpeg png process took {} seconds".format(end - start))
|
|
|
|
else:
|
|
raise ValueError("format must be either PNG or JPEG")
|
|
|
|
return (folder_out,)
|
|
|
|
|
|
class VideoFramesFolderToImageStack(metaclass=ABCWidgetMetaclass):
|
|
"""Use cv2 to open the video and save the videos individual frames"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
from addict import Dict
|
|
ret = Dict()
|
|
ret.required = Dict()
|
|
ret.required.folder_in = ("STRING", {"default": "full_path_to_folder"})
|
|
ret.required.use_subfolders = ([True, False], {"default": False})
|
|
ret.required.by_date_or_name = (["DATE", "NAME"], {"default": "DATE"})
|
|
ret.required.idx_slice_start = ("INT", {"min": 0, "max": 99999, "step": 1, "default": 0})
|
|
ret.required.idx_slice_stop = ("INT", {"min": 0, "max": 99999, "step": 1, "default": 1})
|
|
ret.required.slice_idx_step = ("INT", {"min": 0, "max": 9999, "step": 1, "default": 1})
|
|
ret.required.idx = ("INT", {"min": 0, "max": 99999, "step": 1, "default": 0})
|
|
ret.required.use_float16 = ([True, False], {"default": False})
|
|
|
|
return ret
|
|
|
|
CATEGORY = "video"
|
|
RETURN_TYPES = ("IMAGE", "LIST",)
|
|
RETURN_NAMES = ("image_stack", "[img_fname]",)
|
|
|
|
FUNCTION = "handler"
|
|
|
|
@torch.inference_mode()
|
|
def handler(self, folder_in,
|
|
use_subfolders,
|
|
by_date_or_name,
|
|
idx_slice_start,
|
|
idx_slice_stop,
|
|
slice_idx_step,
|
|
idx,
|
|
use_float16
|
|
):
|
|
"""
|
|
Return slices of time from the video, sequences of images, defined by the start, stop and step
|
|
- torch tensors (T,H,W,C)
|
|
"""
|
|
|
|
import glob
|
|
idx = idx * slice_idx_step
|
|
|
|
folder_in_path = folder_in
|
|
image_list = []
|
|
image_fname_list = []
|
|
|
|
p = os.path.abspath(folder_in_path)
|
|
# List and sort all PNG or JPG files in the folder
|
|
if not use_subfolders and by_date_or_name == "NAME":
|
|
file_fps = sorted(
|
|
glob.glob(os.path.join(p, '*.png')) + glob.glob(os.path.join(p, '*.jpg')))
|
|
elif not use_subfolders and by_date_or_name == "DATE":
|
|
file_fps = sorted(
|
|
glob.glob(os.path.join(p, '*.png')) + glob.glob(os.path.join(p, '*.jpg')),
|
|
key=os.path.getmtime)
|
|
elif use_subfolders and by_date_or_name == "NAME":
|
|
file_fps = sorted(
|
|
glob.glob(os.path.join(p, '**/*.png'), recursive=True) +
|
|
glob.glob(os.path.join(p, '**/*.jpg'), recursive=True))
|
|
elif use_subfolders and by_date_or_name == "DATE":
|
|
file_fps = sorted(
|
|
glob.glob(os.path.join(p, '**/*.png'), recursive=True) +
|
|
glob.glob(os.path.join(p, '**/*.jpg'), recursive=True),
|
|
key=os.path.getmtime)
|
|
|
|
# Ensure the slice doesn't go beyond the end of the file list
|
|
|
|
if idx_slice_stop + idx > len(file_fps):
|
|
idx_slice_stop = len(file_fps) - idx
|
|
|
|
slice_size = idx_slice_stop - idx_slice_start
|
|
# for i in range(idx_slice_start + idx, idx_slice_stop + idx):
|
|
for fn in file_fps[idx_slice_start + idx:idx_slice_start + idx + slice_size]:
|
|
image = cv2.imread(fn)
|
|
image_fname_list.append(os.path.basename(fn))
|
|
if image is not None:
|
|
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) # converting to RGB format
|
|
image_list.append(image)
|
|
|
|
# Convert list of images to tensor (T,H,W,C)
|
|
if use_float16:
|
|
tensor = torch.stack([torch.tensor(img, dtype=torch.float16) for img in image_list])
|
|
else:
|
|
tensor = torch.stack([torch.tensor(img, dtype=torch.float32) for img in image_list])
|
|
del image_list
|
|
tensor = tensor / 255.0
|
|
return (tensor, image_fname_list,)
|
|
|
|
|
|
class SmoothStackTemporalByDistance2(ABCABCVideoFolderToImage, metaclass=ABCWidgetMetaclass):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required":
|
|
{
|
|
"image_stack": ("IMAGE",),
|
|
"top_thresh_normed": ("FLOAT", {"min": 0.0, "max": 100000.0, "step": 0.1, "default": 1.0}),
|
|
"bottom_thresh_normed": ("FLOAT", {"min": 0.0, "max": 100000.0, "step": 0.1, "default": 1.0}),
|
|
"max_iterations": ("INT", {"min": 0, "max": 1000, "step": 1, "default": 1}),
|
|
"target_deviation": ("FLOAT", {"min": 0.0, "max": 100000.0, "step": 0.01, "default": 0.0}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
|
|
def handler(self, image_stack, top_thresh_normed=0.9, bottom_thresh_normed=0.10, max_iterations=1,
|
|
target_deviation=0.0):
|
|
import torch
|
|
# Code for printing memory size and available GPU memory stays the same
|
|
|
|
image_stack = image_stack.to(torch.float16)
|
|
info = ""
|
|
iteration = 0
|
|
while True:
|
|
# Compute the first-order differences (gradient) along the time dimension (0)
|
|
first_order_diff = torch.gradient(image_stack, dim=0)[0]
|
|
|
|
# Compute the Euclidean norm of the first-order differences
|
|
euclidean_diff_first_order = torch.norm(first_order_diff.view(first_order_diff.shape[0], -1), dim=1).clone()
|
|
|
|
euclidean_diff_first_order_range = euclidean_diff_first_order.max() - euclidean_diff_first_order.min()
|
|
euclidean_diff_first_order_normed = (euclidean_diff_first_order - euclidean_diff_first_order.min()) / \
|
|
euclidean_diff_first_order_range
|
|
|
|
euclidean_diff_first_order_normed = euclidean_diff_first_order_normed[:-1]
|
|
|
|
start_mean = euclidean_diff_first_order.mean()
|
|
|
|
start_std = euclidean_diff_first_order.std()
|
|
|
|
if start_std <= target_deviation or iteration >= max_iterations:
|
|
break
|
|
|
|
# Identify the frames with differences greater than the top threshold, excluding the last index
|
|
outlier_indices_top = torch.where(euclidean_diff_first_order_normed > top_thresh_normed)[0]
|
|
|
|
# Interpolate frames where the differences are greater than the top threshold
|
|
interpolated_frames_top = 0.5 * (image_stack[outlier_indices_top] + image_stack[outlier_indices_top + 1])
|
|
image_stack[outlier_indices_top + 1] = interpolated_frames_top
|
|
|
|
# Identify the frames with differences less than the bottom threshold, excluding the last index
|
|
outlier_indices_bottom = torch.where(euclidean_diff_first_order_normed < bottom_thresh_normed)[0]
|
|
|
|
# Create a mask to delete frames where the differences are less than the bottom threshold
|
|
mask = torch.ones(image_stack.size(0), dtype=torch.bool)
|
|
mask[outlier_indices_bottom + 1] = False
|
|
image_stack = image_stack[mask]
|
|
# now remove the first and last frames
|
|
# image_stack = image_stack[1:-1]
|
|
# print the shape and stats for image_stack
|
|
|
|
# relate how many frames where changed via top threshold and bottom threshold
|
|
# also report on the mean and std of the euclidean_diff_first_order
|
|
|
|
info += f"Number of frames changed via top threshold: {len(outlier_indices_top)}\n " \
|
|
f"Number of frames changed via bottom threshold: {len(outlier_indices_bottom)}\n " \
|
|
f"Mean of euclidean_diff_first_order: {start_mean}\n " \
|
|
f"Std of euclidean_diff_first_order: {start_std}\n " \
|
|
f"\n\n"
|
|
print(info)
|
|
iteration += 1
|
|
|
|
if image_stack.shape[0] == 0:
|
|
print("image_stack is empty, returning None")
|
|
return (None, None,)
|
|
print(
|
|
f"image_stack shape: {image_stack.shape}, image_stack stats: {image_stack.min()}, {image_stack.max()}")
|
|
|
|
return (image_stack.to(torch.float32), info,)
|
|
|
|
|
|
class SmoothStackTemporalByDistance(ABCABCVideoFolderToImage, metaclass=ABCWidgetMetaclass):
|
|
"""
|
|
Smooth the temporal dimension of the image stack by calculating the distance between each frame
|
|
and the next frame and then averaging the frames that are within a certain distance of each other.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required":
|
|
{
|
|
"image_stack": ("IMAGE",),
|
|
"top_thresh_normed": ("FLOAT", {"min": 0.0, "max": 100000.0, "step": 0.1, "default": 1.0}),
|
|
"bottom_thresh_normed": ("FLOAT", {"min": 0.0, "max": 100000.0, "step": 0.1, "default": 1.0}),
|
|
"iterations": ("INT", {"min": 0, "max": 1000, "step": 1, "default": 1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
|
|
def handler(self, image_stack,
|
|
top_thresh_normed=0.9, bottom_thresh_normed=0.10, iterations=1):
|
|
import torch
|
|
# calculate and print the total size in memory in bytes of the tensor
|
|
print(
|
|
f"Size of tensor in memory in GB: {image_stack.element_size() * image_stack.nelement() / (1024 ** 3)} GB, dtype: {image_stack.dtype}")
|
|
# print the available memory on the GPU in GB
|
|
print(f"Available memory on GPU in GB: {torch.cuda.get_device_properties(0).total_memory / (1024 ** 3)} GB")
|
|
|
|
# image_stack = image_stack.clone()
|
|
image_stack = image_stack.to(torch.float16)
|
|
info = ""
|
|
for _ in range(iterations):
|
|
# Compute the first-order differences (gradient) along the time dimension (0)
|
|
first_order_diff = torch.gradient(image_stack, dim=0)[0]
|
|
|
|
# Compute the Euclidean norm of the first-order differences
|
|
euclidean_diff_first_order = torch.norm(first_order_diff.view(first_order_diff.shape[0], -1), dim=1).clone()
|
|
|
|
euclidean_diff_first_order_range = euclidean_diff_first_order.max() - euclidean_diff_first_order.min()
|
|
euclidean_diff_first_order_normed = (euclidean_diff_first_order - euclidean_diff_first_order.min()) / \
|
|
euclidean_diff_first_order_range
|
|
|
|
euclidean_diff_first_order_normed = euclidean_diff_first_order_normed[:-1]
|
|
|
|
start_mean = euclidean_diff_first_order.mean()
|
|
start_std = euclidean_diff_first_order.std()
|
|
|
|
# Identify the frames with differences greater than the top threshold, excluding the last index
|
|
outlier_indices_top = torch.where(euclidean_diff_first_order_normed > top_thresh_normed)[0]
|
|
|
|
# Interpolate frames where the differences are greater than the top threshold
|
|
interpolated_frames_top = 0.5 * (image_stack[outlier_indices_top] + image_stack[outlier_indices_top + 1])
|
|
image_stack[outlier_indices_top + 1] = interpolated_frames_top
|
|
|
|
# Identify the frames with differences less than the bottom threshold, excluding the last index
|
|
outlier_indices_bottom = torch.where(euclidean_diff_first_order_normed < bottom_thresh_normed)[0]
|
|
|
|
# Create a mask to delete frames where the differences are less than the bottom threshold
|
|
mask = torch.ones(image_stack.size(0), dtype=torch.bool)
|
|
mask[outlier_indices_bottom + 1] = False
|
|
image_stack = image_stack[mask]
|
|
# now remove the first and last frames
|
|
# image_stack = image_stack[1:-1]
|
|
# print the shape and stats for image_stack
|
|
|
|
# relate how many frames where changed via top threshold and bottom threshold
|
|
# also report on the mean and std of the euclidean_diff_first_order
|
|
info += f"Number of frames changed via top threshold: {len(outlier_indices_top)}\n " \
|
|
f"Number of frames changed via bottom threshold: {len(outlier_indices_bottom)}\n " \
|
|
f"Mean of euclidean_diff_first_order: {start_mean}\n " \
|
|
f"Std of euclidean_diff_first_order: {start_std}\n " \
|
|
f"\n\n"
|
|
print(info)
|
|
|
|
if image_stack.shape[0] == 0:
|
|
print("image_stack is empty, returning None")
|
|
return (None, None,)
|
|
print(
|
|
f"image_stack shape: {image_stack.shape}, image_stack stats: {image_stack.min()}, {image_stack.max()}")
|
|
|
|
return (image_stack.to(torch.float32), info,)
|
|
|
|
|
|
class GetImageStackStatistics(ABCABCVideoFolderToImage, metaclass=ABCWidgetMetaclass):
|
|
"""Get statistics of an image stack"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required":
|
|
{
|
|
"image_stack": ("IMAGE",),
|
|
"method": (["SSIM", "MSE"], {"default": "SSIM"}),
|
|
},
|
|
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "IMAGE",)
|
|
|
|
def handler(self, image_stack, method):
|
|
import numpy as np
|
|
import torch
|
|
# if the dtype is not float32
|
|
if image_stack.dtype != torch.float32:
|
|
image_stack = image_stack.to(torch.float32)
|
|
|
|
image_stack_BCWH = image_stack.permute(0, 3, 1, 2)
|
|
image_stack_BCWH_DIFF = image_stack_BCWH[1:] - image_stack_BCWH[:-1]
|
|
|
|
# we want to calculate the mean squared error
|
|
# of the image stack
|
|
# so we need to calculate the difference between each image in the stack and the previous image
|
|
# we have that in image_stack_BCWH_DIFF
|
|
# now we need to calculate the mean squared error of each image in the stack
|
|
# which will give us a tensor of shape (B)
|
|
# where B is the number of images in the stack
|
|
|
|
if method == "MSE":
|
|
stack_mse_B = torch.mean(image_stack_BCWH_DIFF ** 2, dim=(1, 2, 3))
|
|
# find the max possible value of the mean squared error
|
|
|
|
else:
|
|
# not implimented
|
|
raise NotImplementedError
|
|
|
|
image_stack_diff_len = len(stack_mse_B)
|
|
required_image_info_width = 512 + image_stack_diff_len
|
|
|
|
mean = torch.mean(stack_mse_B)
|
|
std = torch.std(stack_mse_B)
|
|
median = torch.median(stack_mse_B)
|
|
max_mse = torch.max(stack_mse_B)
|
|
min_mse = torch.min(stack_mse_B)
|
|
|
|
out_str = f"mean: {mean}\n std: {std}\n median: {median}\n max: {max_mse}\n min: {min_mse}\n len: {image_stack_diff_len}"
|
|
|
|
# create the histogram
|
|
hist_values, _ = np.histogram(stack_mse_B.numpy(), bins=min(100, image_stack_diff_len))
|
|
hist_values = hist_values / hist_values.max() # Normalize to [0, 1]
|
|
|
|
# Create an image tensor with shape (1, 512, required_width, 3) filled with zeros
|
|
out_hist_image = torch.zeros((1, 512, required_image_info_width, 3))
|
|
|
|
# Determine the scaling factors for width and height for the histogram
|
|
scale_width = 512 // len(hist_values)
|
|
scale_height = out_hist_image.shape[1]
|
|
|
|
# Fill the image tensor with the histogram
|
|
for i, value in enumerate(hist_values):
|
|
height = 512 - int(value * scale_height)
|
|
start_x = int(i * scale_width)
|
|
end_x = int((i + 1) * scale_width)
|
|
out_hist_image[0, height:, start_x:end_x - 1, :] = 1 # Color as white
|
|
|
|
# repeat each element of the stack_mse_B tensor by scale_width
|
|
stack_mse_B = stack_mse_B.repeat_interleave(scale_width)
|
|
|
|
# Create a plot for the difference for each frame pair
|
|
plot_offset = 512
|
|
for i, value in enumerate(stack_mse_B):
|
|
height = 512 - int(value * (scale_height / .5)) # Normalize by the max value
|
|
start_x = plot_offset + i
|
|
end_x = plot_offset + i + 1
|
|
out_hist_image[0, height:, start_x:end_x, :] = (1 * height) / 512 # Color as white
|
|
|
|
return (out_str, out_hist_image,)
|
|
|
|
|
|
# @profile
|
|
class VideoFileToImageStack(ABCABCVideoFileToImage, metaclass=ABCWidgetMetaclass):
|
|
"""Use cv2 to open a video and save the videos individual frames"""
|
|
RETURN_TYPES = ("IMAGE", "FUNC",)
|
|
RETURN_NAMES = ("image_stack", "FUNC(**kwargs)",)
|
|
|
|
def handler(self, video_in, idx_slice_start, idx_slice_stop, slice_idx_step, idx, func_only=False):
|
|
"""
|
|
Return slices of time from the video, sequences of images, defined by the start, stop and step
|
|
- torch tensors (T,H,W,C)
|
|
"""
|
|
import sys
|
|
import gc
|
|
# print memory usage
|
|
print(f"Memory usage: {sys.getsizeof(self)} bytes")
|
|
|
|
if func_only:
|
|
return (None, VideoFileToImageStack.handler,)
|
|
|
|
image_list = []
|
|
video_in_fp = video_in
|
|
vidcap = cv2.VideoCapture(video_in_fp)
|
|
max_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
|
|
|
|
if not vidcap.isOpened():
|
|
print("Error: Unable to open video file")
|
|
return
|
|
|
|
start = (idx * slice_idx_step) + idx_slice_start
|
|
stop = (idx * slice_idx_step) + idx_slice_stop
|
|
slice_length = stop - start
|
|
|
|
if (start + slice_length) > max_frames:
|
|
raise ValueError(f"idx_slice_stop ({idx_slice_stop}) is greater than the max frames ({max_frames})")
|
|
|
|
for frame_count in range(start, start + slice_length):
|
|
# Seek to frame_count frame
|
|
vidcap.set(cv2.CAP_PROP_POS_FRAMES, frame_count)
|
|
|
|
success, image = vidcap.read()
|
|
if not success:
|
|
print("Error: Unable to read frame from video file")
|
|
break
|
|
# image_list.append(image)
|
|
|
|
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
image_list.append(image)
|
|
print(f"Memory usage: {sys.getsizeof(self)} bytes")
|
|
|
|
vidcap.release()
|
|
# Convert list of images to tensor (T,H,W,C)
|
|
try:
|
|
tensor = torch.stack([torch.tensor(img, dtype=torch.float32) for img in image_list])
|
|
except RuntimeError as e:
|
|
print("Error: Unable to convert image list to tensor")
|
|
raise e
|
|
|
|
tensor = tensor / 255.0
|
|
del image_list
|
|
del vidcap
|
|
del image
|
|
gc.collect()
|
|
# print memory usage
|
|
print(f"Memory usage: {sys.getsizeof(self)} bytes")
|
|
return (tensor, VideoFileToImageStack.handler,)
|
|
|
|
|
|
class ImageStackToVideoFramesFolder(metaclass=ABCWidgetMetaclass):
|
|
"""Really just saves the image stack to a folder"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(self):
|
|
return {"required":
|
|
{
|
|
"image_stack": ("IMAGE",),
|
|
"folder_out": ("STRING", {"multiline": False, "default": "video01"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
CATEGORY = "video"
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "handler"
|
|
|
|
def handler(self, **kwargs):
|
|
import tempfile
|
|
import uuid
|
|
import shutil
|
|
|
|
image_stack = kwargs["image_stack"]
|
|
folder_out = kwargs["folder_out"]
|
|
folder_out_path = video_dir
|
|
|
|
# full path to the output folder
|
|
folder_out_path = os.path.join(video_dir, folder_out)
|
|
|
|
# Create the output folder if it doesn't exist
|
|
if not os.path.exists(folder_out_path):
|
|
os.makedirs(folder_out_path)
|
|
|
|
# move any current files in the output folder to a random folder inside the video folder
|
|
random_folder = os.path.join(folder_out_path, str(uuid.uuid4()))
|
|
os.makedirs(random_folder)
|
|
for file in os.listdir(folder_out_path):
|
|
file_path = os.path.join(folder_out_path, file)
|
|
if os.path.isfile(file_path):
|
|
shutil.move(file_path, random_folder)
|
|
|
|
# Save each image in the image stack to the output folder
|
|
for i in range(image_stack.shape[0]):
|
|
image = image_stack[i].numpy()
|
|
# make sure the image is in the correct color format
|
|
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
|
|
# convert to 0-255
|
|
image = image * 255
|
|
# use zfill to get a good ordering of the frames
|
|
file_name = os.path.join(folder_out_path, "frame" + str(i).zfill(6) + ".jpg")
|
|
# use the best compression nearly lossless
|
|
cv2.imwrite(file_name, image, [cv2.IMWRITE_JPEG_QUALITY, 99])
|
|
|
|
image_stack = None
|
|
del image_stack
|
|
return (folder_out,)
|
|
|
|
|
|
class TemporalSpatialSmoothing(metaclass=ABCWidgetMetaclass):
|
|
"""Smooths the image stack in time and space"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(self):
|
|
return {"required":
|
|
{
|
|
"image_stack": ("IMAGE",),
|
|
"spatial_kernel_size": ("INT", {"default": 3, "min": 1, "max": 50}),
|
|
"temporal_kernel_size": ("INT", {"default": 3, "min": 1, "max": 50}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "video"
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "handler"
|
|
|
|
@staticmethod
|
|
def gaussian_kernel_1d(kernel_size=3, sigma=1.0):
|
|
import numpy as np
|
|
if kernel_size % 2 == 0:
|
|
kernel_size += 1
|
|
x_coord = torch.arange(-kernel_size // 2 + 1., kernel_size // 2 + 1.).float()
|
|
gaussian_kernel = (1. / (sigma * torch.sqrt(torch.tensor(2. * np.pi)))) * torch.exp(
|
|
-x_coord ** 2. / (2 * sigma ** 2))
|
|
gaussian_kernel = gaussian_kernel / torch.sum(gaussian_kernel)
|
|
return gaussian_kernel
|
|
|
|
def handler(self, **kwargs):
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
image_stack = kwargs["image_stack"]
|
|
spatial_kernel_size = kwargs["spatial_kernel_size"]
|
|
temporal_kernel_size = kwargs["temporal_kernel_size"]
|
|
spatial_sigma = kwargs.get("spatial_sigma", 1.0)
|
|
temporal_sigma = kwargs.get("temporal_sigma", 1.0)
|
|
|
|
# make sure image_stack is a torch tensor and has the right dimensions
|
|
if not torch.is_tensor(image_stack):
|
|
image_stack = torch.tensor(image_stack, dtype=torch.float32)
|
|
|
|
if len(image_stack.shape) < 5: # if it doesn't have a channel dimension
|
|
image_stack = image_stack.unsqueeze(0) # add a batch dimension
|
|
image_stack = image_stack.permute(0, 4, 1, 2,
|
|
3) # rearrange dimensions to: (batch, channel, depth, height, width)
|
|
|
|
# create 1D Gaussian kernels for each dimension
|
|
spatial_kernel = self.gaussian_kernel_1d(spatial_kernel_size, spatial_sigma)
|
|
temporal_kernel = self.gaussian_kernel_1d(temporal_kernel_size, temporal_sigma)
|
|
|
|
# apply the convolution along each dimension separately
|
|
for dim, kernel in enumerate([temporal_kernel, spatial_kernel, spatial_kernel]):
|
|
padding_size = kernel.shape[0] // 2
|
|
# define the padding differently for each dimension
|
|
if dim == 0: # depth (temporal)
|
|
padding = nn.ReplicationPad3d((0, 0, 0, 0, padding_size, padding_size))
|
|
elif dim == 1: # height
|
|
padding = nn.ReplicationPad3d((0, 0, padding_size, padding_size, 0, 0))
|
|
elif dim == 2: # width
|
|
padding = nn.ReplicationPad3d((padding_size, padding_size, 0, 0, 0, 0))
|
|
image_stack = padding(image_stack)
|
|
if dim == 0: # depth/temporal
|
|
conv = nn.Conv3d(in_channels=3, out_channels=3, kernel_size=(kernel.shape[0], 1, 1), groups=3,
|
|
bias=False)
|
|
reshaped_kernel = kernel[None, None, :, None, None]
|
|
elif dim == 1: # height
|
|
conv = nn.Conv3d(in_channels=3, out_channels=3, kernel_size=(1, kernel.shape[0], 1), groups=3,
|
|
bias=False)
|
|
reshaped_kernel = kernel[None, None, None, :, None]
|
|
elif dim == 2: # width
|
|
conv = nn.Conv3d(in_channels=3, out_channels=3, kernel_size=(1, 1, kernel.shape[0]), groups=3,
|
|
bias=False)
|
|
reshaped_kernel = kernel[None, None, None, None, :]
|
|
conv.weight = nn.Parameter(reshaped_kernel.repeat(3, 1, 1, 1, 1), requires_grad=False)
|
|
image_stack = conv(image_stack)
|
|
|
|
# remove the batch dimension and rearrange dimensions back to: (depth, height, width, channel)
|
|
image_stack = image_stack.squeeze(0).permute(1, 2, 3, 0)
|
|
|
|
return (image_stack,)
|
|
|
|
|
|
class TemporalOpticalFlowSmoothing(metaclass=ABCWidgetMetaclass):
|
|
"""Smooths the image stack in time using RAFT for optical flow estimation"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required":
|
|
{
|
|
"image_stack": ("IMAGE",),
|
|
"k": ("INT", {"default": 3, "min": 1, "max": 120}),
|
|
"device": ("STRING", ["cpu", "cuda"]),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "video"
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "handler"
|
|
|
|
@staticmethod
|
|
def make_grid(flow):
|
|
# Assumes flow is in shape Bx2xHxW
|
|
B, _, H, W = flow.size()
|
|
grid_y, grid_x = torch.meshgrid(torch.linspace(-1, 1, W), torch.linspace(-1, 1, H))
|
|
grid = torch.stack((grid_y.t(), grid_x.t()), 2).unsqueeze(0)
|
|
grid = grid.repeat(B, 1, 1, 1) # Make it BxHxWx2
|
|
grid = grid.permute(0, 3, 1, 2) # Convert to Bx2xHxW
|
|
grid = grid.type_as(flow) # Make sure the data type is the same
|
|
return grid + flow
|
|
|
|
@staticmethod
|
|
def worker(args):
|
|
import torch.nn.functional as F
|
|
k, image_stack, device = args
|
|
weights = Raft_Large_Weights.DEFAULT
|
|
model = raft_large(weights=weights, progress=False).to(device)
|
|
model = model.eval()
|
|
|
|
img1 = image_stack[0].unsqueeze(0)
|
|
img2 = image_stack[1].unsqueeze(0)
|
|
|
|
# Remove batch dimension if present
|
|
while len(img1.shape) > 4:
|
|
print(f"Batch dimension detected. for image shape {img1.shape}")
|
|
img1 = img1.squeeze(0)
|
|
img2 = img2.squeeze(0)
|
|
|
|
# Permute tensor dimensions to match expected input for model
|
|
img1 = img1.permute(0, 3, 1, 2)
|
|
img2 = img2.permute(0, 3, 1, 2)
|
|
|
|
# Move images to the specified device
|
|
img1, img2 = img1.to(device), img2.to(device)
|
|
|
|
flow = model(img1, img2)[-1]
|
|
|
|
# Create intermediary frames
|
|
frames = []
|
|
for i in range(k):
|
|
factor = i / k
|
|
# factor = factor * 2 - 1
|
|
|
|
# normalize flow
|
|
n_flow = flow.clone()
|
|
print(n_flow.min(), n_flow.max(), n_flow.dtype)
|
|
# n_flow[:,0,:,:] = n_flow[:,0,:,:] / (n_flow[:,0,:,:].max() - n_flow[:,0,:,:].min())
|
|
# n_flow[:,1,:,:] = n_flow[:,1,:,:] / (n_flow[:,1,:,:].max() - n_flow[:,1,:,:].min())
|
|
# print(n_flow.min(), n_flow.max(), n_flow.dtype)
|
|
# n_flow[:, 0, :, :] = n_flow[:, 0, :, :] - n_flow[:, 0, :, :].min()
|
|
# n_flow[:, 1, :, :] = n_flow[:, 1, :, :] - n_flow[:, 1, :, :].min()
|
|
# print (n_flow.min(), n_flow.max(),n_flow.dtype)
|
|
n_flow = n_flow / n_flow.shape[2]
|
|
print(n_flow.min(), n_flow.max(), n_flow.dtype)
|
|
|
|
flow_interpolated = n_flow * factor
|
|
grid = TemporalOpticalFlowSmoothing.make_grid(flow_interpolated)
|
|
grid = grid.permute(0, 2, 3, 1)
|
|
warped_image = F.grid_sample(img1, grid, mode='bicubic', padding_mode='zeros')
|
|
warped_image = warped_image.permute(0, 2, 3, 1).detach().cpu() # Convert it back to BxHxWxC
|
|
frames.append(warped_image)
|
|
|
|
return frames
|
|
|
|
@staticmethod
|
|
def calculate_optical_flow(image_stack, device, k):
|
|
# from torch.multiprocessing import Pool, set_start_method
|
|
|
|
# set_start_method('spawn') # set the start method to spawn
|
|
# with Pool(2) as pool:
|
|
# # create the tasks in an intellegent way, detaching each stack, with each stack being only the frames needed
|
|
# tasks = []
|
|
# for t in range(image_stack.shape[0] - 1):
|
|
# needed_frames = [t, t + 1] # the frames needed for this task
|
|
# stack = image_stack[needed_frames, :, :, :] # the stack of frames needed for this task
|
|
# stack = stack.clone()
|
|
# tasks.append((2, stack.detach().cpu(), device)) # add the task to the list of tasks
|
|
# flows = pool.map(TemporalOpticalFlowSmoothing.worker, tasks)
|
|
# # print(f"flows: {flows}")
|
|
# # print(f"flows[0]: {flows[0]}")
|
|
|
|
##map(TemporalOpticalFlowSmoothing.worker, tasks)
|
|
|
|
flows = []
|
|
for i in range(image_stack.shape[0] - 1):
|
|
needed_frames = [i, i + 1] # the frames needed for this task
|
|
keys = image_stack[needed_frames, :, :, :]
|
|
flow = TemporalOpticalFlowSmoothing.worker((k, keys, device))
|
|
|
|
flows.append(keys[0].cpu())
|
|
for f in flow:
|
|
flows.append(f[0].cpu())
|
|
# flows.append(keys[1].cpu())
|
|
|
|
return flows
|
|
|
|
def handler(self, **kwargs):
|
|
from copy import deepcopy
|
|
# image stacks are actually (b/t,w,h,c) so we need to rearrange the dimensions
|
|
k = kwargs["k"]
|
|
image_stack = deepcopy(kwargs["image_stack"])
|
|
device = kwargs.get("device", "cpu")
|
|
|
|
image_stack = image_stack.permute(0, 2, 1, 3)
|
|
|
|
# make sure image_stack is a torch tensor and has the right dimensions
|
|
if not torch.is_tensor(image_stack):
|
|
image_stack = torch.tensor(image_stack, dtype=torch.float32)
|
|
|
|
if len(image_stack.shape) < 4: # if it doesn't have a channel dimension
|
|
image_stack = image_stack.unsqueeze(0) # add a batch dimension
|
|
|
|
# Move image stack to the specified device
|
|
image_stack = image_stack.to(device)
|
|
|
|
# Calculate optical flow
|
|
flows = self.calculate_optical_flow(image_stack, device, k)
|
|
stack = torch.stack(flows, dim=0)
|
|
# Convert optical flows to images for visualization
|
|
# flow_images = [flow_to_image(flow) for flow in flows]
|
|
# change back to (b or t,w,h,c)
|
|
|
|
stack = stack.permute(0, 2, 1, 3)
|
|
return (stack,)
|
|
|
|
|
|
class MovingCircle(metaclass=ABCWidgetMetaclass):
|
|
# VideoWidget class definition
|
|
|
|
def calculate_start(self, x1, y1, width, height):
|
|
# Convert to 0-1 range
|
|
start_x = x1 * width
|
|
start_y = y1 * height
|
|
return (start_x, start_y)
|
|
|
|
def calculate_end(self, x2, y2, width, height):
|
|
# Convert to 0-1 range
|
|
end_x = x2 * width
|
|
end_y = y2 * height
|
|
return (end_x, end_y)
|
|
|
|
def generate_frames(self, num_frames, width, height):
|
|
# Match tensor shape from example
|
|
frames = torch.zeros(num_frames, width, height, 3)
|
|
return frames
|
|
|
|
def draw_circle_or_oval(self, kp, width, height, radius, aspect_ratio, color=(255, 255, 255)):
|
|
import torch.nn.functional as F
|
|
# Initialize black background
|
|
img = torch.zeros((3, height, width), dtype=torch.uint8)
|
|
|
|
# Draw circle as keypoints
|
|
keypoints = torch.tensor([[kp[0], kp[1]]]).unsqueeze(0)
|
|
img = draw_keypoints(img, keypoints, colors=color, radius=radius)
|
|
|
|
# convert to (w,h,c)
|
|
# img = img.permute(1, 2, 0).unsqueeze(0)
|
|
img = img.unsqueeze(0)
|
|
|
|
# Calculate new size
|
|
if aspect_ratio < 1:
|
|
new_size = (height, round(width / aspect_ratio))
|
|
else:
|
|
new_size = (round(height * aspect_ratio), width)
|
|
|
|
# Stretch the image
|
|
img_stretched = F.interpolate(img, size=new_size, mode='bilinear', align_corners=False)
|
|
|
|
# Shrink it back
|
|
img_shrunk = F.interpolate(img_stretched, size=(height, width), mode='bilinear', align_corners=False)
|
|
|
|
# Convert back to float values
|
|
|
|
img_shrunk = img_shrunk.float() / 255
|
|
|
|
img_shrunk = img_shrunk.permute(0, 2, 3, 1)
|
|
return img_shrunk
|
|
|
|
def draw_circle(self, kp, width, height, radius, color=(255, 255, 255)):
|
|
# Initialize black background
|
|
img = torch.zeros((3, height, width), dtype=torch.uint8)
|
|
# Draw circle as keypoints
|
|
keypoints = torch.tensor([[kp[0], kp[1]]]).unsqueeze(0)
|
|
img = draw_keypoints(img, keypoints, colors=color, radius=radius)
|
|
# convert to (w,h,c)
|
|
img = img.permute(1, 2, 0).unsqueeze(0)
|
|
# convert back to float values
|
|
img = img.float() / 255
|
|
|
|
return img
|
|
|
|
# Metadata for documentation
|
|
RETURN_TYPES = ("IMAGE",)
|
|
CATEGORY = "ETK/video"
|
|
FUNCTION = "handler"
|
|
|
|
# Match parent class method
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
# Validate types, ranges
|
|
return {"required":
|
|
{
|
|
"x1": ("FLOAT", {"min": 0, "max": 1, "step": 0.1}),
|
|
"y1": ("FLOAT", {"min": 0, "max": 1, "step": 0.1}),
|
|
"x2": ("FLOAT", {"min": 0, "max": 1, "step": 0.1}),
|
|
"y2": ("FLOAT", {"min": 0, "max": 1, "step": 0.1}),
|
|
"num_frames": ("INT", {"min": 1}),
|
|
"width": ("INT", {"min": 1}),
|
|
"height": ("INT", {"min": 1}),
|
|
"radius": ("INT", {"min": 1}),
|
|
"color": ("STRING", {"default": "255 255 255"}),
|
|
"aspect_ratio": ("FLOAT", {"min": 0.1, "max": 10})
|
|
|
|
}}
|
|
|
|
def handler(self, x1, y1, x2, y2, num_frames, width, height, radius, color="255 255 255", aspect_ratio=1):
|
|
import torch
|
|
# Avoid errors on bad values
|
|
assert 0 <= x1 <= 1
|
|
assert 0 <= y1 <= 1
|
|
|
|
# Split calculations
|
|
start = torch.tensor(self.calculate_start(x1, y1, width, height))
|
|
|
|
# Split calculations
|
|
end = torch.tensor(self.calculate_end(x2, y2, width, height))
|
|
|
|
# color is supposed to be a tuple
|
|
color = tuple(map(lambda x: int(x), color.split()))
|
|
# Initialize for compositing
|
|
|
|
frames = []
|
|
for i in range(num_frames):
|
|
# Leverage PyTorch
|
|
cur_pos = start + (end - start) * i / num_frames
|
|
circle = self.draw_circle_or_oval(cur_pos, width, height, radius, aspect_ratio, color)
|
|
# Use broadcasting to overlay
|
|
frames.append(circle)
|
|
|
|
# cat
|
|
frames = torch.cat(frames, dim=0)
|
|
|
|
return (frames,)
|
|
|
|
if __name__ == "__main__":
|
|
import doctest
|
|
|
|
doctest.testmod()
|
|
|
|
|
|
class FindFrameFolders(metaclass=ABCWidgetMetaclass):
|
|
"""Finds all folders directly off a given root path containing sequential image files"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(self):
|
|
return {
|
|
"required": {
|
|
"root_path": ("STRING", {"default": "path to root of project"})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LIST",)
|
|
RETURN_NAMES = ("image_sequence_folders",)
|
|
CATEGORY = "video"
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "handler"
|
|
|
|
def handler(self, **kwargs):
|
|
root_path = kwargs["root_path"]
|
|
image_sequence_folders = []
|
|
|
|
for folder_name in os.listdir(root_path):
|
|
folder_path = os.path.join(root_path, folder_name)
|
|
if os.path.isdir(folder_path):
|
|
image_files = [file for file in os.listdir(folder_path) if file.endswith(('.png', '.jpg', '.jpeg'))]
|
|
if len(image_files) >= 2:
|
|
image_sequence_folders.append(folder_path)
|
|
|
|
return (image_sequence_folders,)
|
|
|
|
|
|
class CombineAudioAndVideoFiles(metaclass=ABCWidgetMetaclass):
|
|
"""Combines lists of audio files and video files using imageio-ffmpeg"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(self):
|
|
return {
|
|
"required": {
|
|
"audio_files": ("LIST", {"default": []}),
|
|
"video_files": ("LIST", {"default": []}),
|
|
"fps": ("FLOAT", {"default": 30, "min": 1, "max": 120}),
|
|
},
|
|
"optional": {
|
|
"move_to_folder": ("STRING", {"default": ""}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LIST",)
|
|
RETURN_NAMES = ("combined_video_paths",)
|
|
CATEGORY = "video"
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "handler"
|
|
|
|
@classmethod
|
|
def has_audio_stream(cls, video_file):
|
|
"""Check if the video file has an audio stream."""
|
|
command = ['ffprobe', '-v', 'error', '-select_streams', 'a', '-show_entries', 'stream=codec_name', '-of',
|
|
'default=nw=1:nk=1', video_file]
|
|
try:
|
|
output = subprocess.check_output(command, stderr=subprocess.STDOUT)
|
|
return bool(output.strip())
|
|
except subprocess.CalledProcessError:
|
|
return False
|
|
|
|
def handler(self, **kwargs):
|
|
from os.path import splitext, join
|
|
import shutil
|
|
audio_files = kwargs["audio_files"]
|
|
video_files = kwargs["video_files"]
|
|
move_to_folder = kwargs.get("move_to_folder", "")
|
|
|
|
combined_video_paths = []
|
|
video_audio_pairs = zip(video_files, audio_files)
|
|
|
|
ffmpeg_exe = ffmpeg.get_ffmpeg_exe()
|
|
if ffmpeg_exe is None:
|
|
raise RuntimeError("FFmpeg could not be found.")
|
|
|
|
for video_file, audio_file in video_audio_pairs:
|
|
|
|
video_out_file_fp = splitext(video_file)[0] + "_v_and_a.mp4"
|
|
input_args = ['-i', video_file, '-i', audio_file]
|
|
|
|
if self.has_audio_stream(video_file):
|
|
# The video has an audio stream, proceed with mixing
|
|
filter_complex_arg = ['-filter_complex', '[0:a][1:a]amix=inputs=2:duration=longest[a]']
|
|
output_args = ['-c:v', 'copy', '-map', '0:v:0', '-map', '[a]', '-c:a', 'aac', '-strict', 'experimental',
|
|
'-y', video_out_file_fp]
|
|
|
|
|
|
else:
|
|
filter_complex_arg = []
|
|
output_args = ['-c:v', 'copy', '-map', '0:v:0', '-map', '1:a:0', '-c:a', 'aac', '-strict',
|
|
'experimental', '-y', video_out_file_fp]
|
|
|
|
command = [ffmpeg_exe] + input_args + filter_complex_arg + output_args
|
|
cmd = ' '.join(command)
|
|
print(cmd)
|
|
os.system(cmd)
|
|
|
|
if move_to_folder:
|
|
destination_path = join(move_to_folder, os.path.basename(video_out_file_fp))
|
|
|
|
if video_out_file_fp != destination_path:
|
|
shutil.move(video_out_file_fp, destination_path)
|
|
|
|
combined_video_paths.append(destination_path)
|
|
else:
|
|
combined_video_paths.append(video_out_file_fp)
|
|
|
|
return (combined_video_paths,)
|
|
|
|
|
|
class CombineFoldersWithFramesToVideoFiles(metaclass=ABCWidgetMetaclass):
|
|
"""Combines lists of folders containing frames of individual video files using imageio and ffmpeg to create video files"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"folders_with_frames": ("LIST", {"default": []}),
|
|
"fps/length": ("FLOAT", {"default": 30, "min": 1, "max": 10000000000}),
|
|
"bitrate": ("STRING", {"default": "20000k"}),
|
|
"codec": (
|
|
["libx264", "libx265", "mpeg4", "vp9", "huffyuv", "flv1", "h264_nvenc", "hevc_nvenc"],
|
|
{"default": "libx264"}
|
|
),
|
|
"file_extension": (
|
|
[".mp4", ".mov", ".avi", ".flv", ".mkv", ".webm"],
|
|
{"default": ".mp4"}
|
|
),
|
|
"fps or length": (["fps", "length"], {"default": "fps"}),
|
|
},
|
|
"optional": {
|
|
"[duration or fps]": ("LIST", {"default": []}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("LIST", "STRING",)
|
|
RETURN_NAMES = ("combined_video_paths",)
|
|
CATEGORY = "video"
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "handler"
|
|
|
|
def handler(self, **kwargs):
|
|
import imageio
|
|
import os
|
|
|
|
folders_with_frames = kwargs["folders_with_frames"]
|
|
fps_length_value = kwargs["fps/length"]
|
|
bitrate = kwargs["bitrate"]
|
|
codec = kwargs["codec"]
|
|
file_extension = kwargs["file_extension"]
|
|
fps_or_length = kwargs["fps or length"]
|
|
duration_or_fps_list = kwargs.get("[duration or fps]", [])
|
|
|
|
combined_video_paths = []
|
|
messages = []
|
|
|
|
for i, folder_path in enumerate(folders_with_frames):
|
|
try:
|
|
out_file = os.path.join(os.path.dirname(folder_path),
|
|
f"{os.path.basename(folder_path)}tmp_{i}{file_extension}")
|
|
frame_files = sorted(
|
|
[f for f in os.listdir(folder_path) if f.lower().endswith(('.jpg', '.jpeg', '.png'))])
|
|
|
|
if not frame_files:
|
|
messages.append(f"No frame files found in folder: {folder_path}")
|
|
continue
|
|
|
|
if fps_or_length == "fps":
|
|
fps = fps_length_value
|
|
if duration_or_fps_list:
|
|
fps = duration_or_fps_list[i]
|
|
|
|
else:
|
|
total_frames = len(frame_files)
|
|
fps = total_frames / fps_length_value if fps_length_value != 0 else 30
|
|
if duration_or_fps_list:
|
|
# calculate the fps based on the duration
|
|
fps = total_frames / duration_or_fps_list[i]
|
|
|
|
writer = imageio.get_writer(out_file, fps=fps, codec=codec, bitrate=bitrate)
|
|
|
|
for frame_file in frame_files:
|
|
frame_path = os.path.join(folder_path, frame_file)
|
|
frame = imageio.imread(frame_path)
|
|
writer.append_data(frame)
|
|
|
|
writer.close()
|
|
combined_video_paths.append(out_file)
|
|
except Exception as e:
|
|
messages.append(f"Error processing folder {folder_path}: {str(e)}")
|
|
|
|
return (combined_video_paths, "\n".join(messages),)
|
|
|
|
|
|
class ImageStackToVideoFile(metaclass=ABCWidgetMetaclass):
|
|
"""Really just saves the image stack to a video file"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(self):
|
|
return {"required":
|
|
{
|
|
"image_stack": ("IMAGE",),
|
|
"video_out": ("STRING", {"multiline": False, "default": "video01.mp4"}),
|
|
"fps or length": (["fps", "length"], {"default": "fps"}),
|
|
"fps/length": ("FLOAT", {"default": 30, "min": 1, "max": 10000000000}),
|
|
"bitrate": ("STRING", {"default": "20000k"}),
|
|
"codec": (
|
|
["libx264", "libx265", "mpeg4", "libvpx-vp9", "huffyuv", "flv1"],
|
|
{"default": "libx264"}
|
|
),
|
|
},
|
|
"optional":
|
|
{
|
|
"audio_file": ("STRING", {"default": "NONE"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("FUNC", "STRING",)
|
|
RETURN_NAMES = ("FUNC", "video_out_path",)
|
|
CATEGORY = "video"
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "handler"
|
|
|
|
def handler(self, **kwargs):
|
|
import os
|
|
import numpy as np
|
|
import uuid
|
|
import shutil
|
|
import imageio
|
|
import audioread
|
|
import cv2
|
|
import torchvision
|
|
|
|
use_torchvision = False
|
|
|
|
# video_dir = kwargs["video_dir"]
|
|
|
|
audio_file = kwargs.get("audio_file")
|
|
if audio_file == "NONE":
|
|
audio_file = None
|
|
|
|
image_stack = kwargs["image_stack"]
|
|
|
|
fps_or_length = kwargs["fps or length"]
|
|
fps_length_value = kwargs["fps/length"]
|
|
bitrate = kwargs["bitrate"]
|
|
codec = kwargs["codec"]
|
|
video_out = kwargs["video_out"]
|
|
|
|
if fps_or_length == "fps":
|
|
fps_length_value = np.array([fps_length_value], dtype=np.float16)[0]
|
|
|
|
video_out_path = video_out
|
|
|
|
if os.path.exists(video_out_path):
|
|
random_folder = os.path.join(os.path.dirname(video_out_path), str(uuid.uuid4()))
|
|
os.makedirs(random_folder)
|
|
shutil.move(video_out_path, random_folder)
|
|
|
|
if audio_file is not None:
|
|
audio_file = audio_file
|
|
if not os.path.exists(audio_file):
|
|
# raise Exception(f"Audio file {audio_file} does not exist")
|
|
audio_length = 0
|
|
else:
|
|
with audioread.audio_open(audio_file) as audio_info:
|
|
audio_length = audio_info.duration
|
|
|
|
frames_length = image_stack.shape[0]
|
|
if fps_or_length == "fps":
|
|
video_length = frames_length / fps_length_value
|
|
else:
|
|
video_length = fps_length_value
|
|
|
|
if fps_or_length == "fps":
|
|
video_fps = fps_length_value
|
|
else:
|
|
video_fps = frames_length / fps_length_value
|
|
|
|
if use_torchvision == False:
|
|
writer = imageio.get_writer(video_out_path, fps=float(video_fps))
|
|
|
|
# convert the image stack to uint8
|
|
image_stack = (image_stack * 255).to(torch.uint8)
|
|
|
|
for frame_data in image_stack:
|
|
frame = np.array(frame_data) # Convert torch tensor to numpy array if needed
|
|
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
writer.append_data(frame)
|
|
|
|
if audio_file is not None:
|
|
difference = video_length - audio_length
|
|
|
|
needed_extra_frames = int(difference * video_fps)
|
|
|
|
if needed_extra_frames > 0:
|
|
for i in range(needed_extra_frames):
|
|
writer.append_data(frame)
|
|
|
|
writer.close()
|
|
|
|
args = [
|
|
"-i", video_out_path,
|
|
"-i", audio_file,
|
|
"-c:v", codec,
|
|
"-b:v", str(bitrate),
|
|
"-c:a", "aac",
|
|
"-strict", "experimental",
|
|
"-y",
|
|
video_out_path.replace(".mp4", "_with_audio.mp4")
|
|
]
|
|
run_ffmpeg_command(args)
|
|
else:
|
|
writer.close()
|
|
|
|
if use_torchvision:
|
|
audio_tensor = None
|
|
if audio_file is not None:
|
|
audio_tensor, aud_sr = torchaudio.load(audio_file)
|
|
fps_frac = fractions.Fraction(float(video_fps))
|
|
torchvision.io.write_video(filename=video_out_path,
|
|
video_array=image_stack * 255,
|
|
fps=fps_frac,
|
|
audio_array=audio_tensor,
|
|
video_codec=codec,
|
|
audio_codec="aac",
|
|
audio_fps=aud_sr
|
|
)
|
|
|
|
return (self, video_out_path,)
|
|
|
|
|
|
class CreateConsistentVideo(metaclass=ABCWidgetMetaclass):
|
|
"""Creates a single video from a folder of videos with a consistent rate of change"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"video_folder": ("STRING", {"default": ""}),
|
|
"rate_of_change_target": ("FLOAT", {"default": 0.5}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "LIST", "STRING",)
|
|
RETURN_NAMES = ("joined_video_path", "individual_video_paths", "error_messages")
|
|
CATEGORY = "video"
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "handler"
|
|
|
|
def handler(self, **kwargs):
|
|
import asyncio
|
|
import os
|
|
|
|
video_folder = kwargs["video_folder"]
|
|
rate_of_change_target = kwargs["rate_of_change_target"]
|
|
|
|
individual_video_paths = []
|
|
error_messages = []
|
|
|
|
async def process_video(file_path):
|
|
try:
|
|
# Calculate visual differences
|
|
difference = calculate_visual_difference(file_path)
|
|
|
|
# Adjust frame rate based on the rate of change target
|
|
adjusted_video_path = adjust_frame_rate(
|
|
video_path=file_path,
|
|
target_rate_of_change=rate_of_change_target,
|
|
visual_difference=difference
|
|
|
|
)
|
|
|
|
return adjusted_video_path
|
|
|
|
except Exception as e:
|
|
error_messages.append(str(e))
|
|
return None
|
|
|
|
# Gather video files
|
|
|
|
video_files = [f for f in os.listdir(video_folder) if
|
|
f.lower().endswith(('.mp4', '.mov', '.avi', '.flv', '.mkv', '.webm'))]
|
|
# now get their full path
|
|
video_files = [os.path.join(video_folder, f) for f in video_files]
|
|
|
|
# Use Asyncio to process videos in parallel
|
|
# There is no current loop so we need to create a NEW one
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
processed_videos = loop.run_until_complete(asyncio.gather(
|
|
*[process_video(file_path) for file_path in video_files]
|
|
))
|
|
|
|
# Combine processed videos into the final video
|
|
joined_video_path = combine_videos(processed_videos, video_folder, video_folder)
|
|
|
|
return (joined_video_path, individual_video_paths, "\n".join(error_messages),)
|
|
|
|
|
|
def calculate_visual_difference(video_path):
|
|
# Open the video file
|
|
cap = cv2.VideoCapture(video_path)
|
|
if not cap.isOpened():
|
|
raise Exception("Could not open the video file")
|
|
|
|
# Initialize variables to hold previous frame
|
|
prev_frame_tensor = None
|
|
total_difference = 0
|
|
|
|
# Iterate through frames
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
|
|
# Convert frame to a PyTorch tensor and normalize
|
|
frame_tensor = torch.from_numpy(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).permute(2, 0, 1).float() / 255
|
|
|
|
# If previous frame exists, calculate the difference using SSIM
|
|
if prev_frame_tensor is not None:
|
|
difference = get_img_diff_torch(prev_frame_tensor, frame_tensor)
|
|
total_difference += difference
|
|
|
|
# Update previous frame
|
|
prev_frame_tensor = frame_tensor
|
|
|
|
# Close the video file
|
|
cap.release()
|
|
|
|
return total_difference
|
|
|
|
|
|
def get_img_diff_torch(image1: torch.Tensor, image2: torch.Tensor, K1=0.01, K2=0.03, L=255):
|
|
# Convert images to grayscale using weighted average of color channels
|
|
gray1 = 0.2989 * image1[0] + 0.5870 * image1[1] + 0.1140 * image1[2]
|
|
gray2 = 0.2989 * image2[0] + 0.5870 * image2[1] + 0.1140 * image2[2]
|
|
|
|
# Calculate mean and variance of input images
|
|
mu1 = torch.mean(gray1)
|
|
mu2 = torch.mean(gray2)
|
|
var1 = torch.var(gray1)
|
|
var2 = torch.var(gray2)
|
|
covar = np.cov(gray1.cpu().numpy().ravel(), gray2.cpu().numpy().ravel())[0][1]
|
|
|
|
# Set constants
|
|
C1 = (K1 * L) ** 2
|
|
C2 = (K2 * L) ** 2
|
|
|
|
# Calculate SSIM
|
|
numerator = (2 * mu1 * mu2 + C1) * (2 * covar + C2)
|
|
denominator = (mu1 ** 2 + mu2 ** 2 + C1) * (var1 + var2 + C2)
|
|
ssim = numerator / denominator
|
|
|
|
# Return difference score
|
|
return 1 - ssim
|
|
|
|
|
|
def adjust_frame_rate(video_path, target_rate_of_change, visual_difference):
|
|
filename, _ = os.path.splitext(os.path.basename(video_path))
|
|
base_dir = os.path.dirname(video_path)
|
|
output_filename = f"{filename}_adjusted.mp4"
|
|
output_filename = os.path.join(base_dir, output_filename)
|
|
|
|
# use the full path
|
|
|
|
cap = cv2.VideoCapture(video_path)
|
|
if not cap.isOpened():
|
|
raise Exception("Could not open the video file")
|
|
|
|
original_fps = int(cap.get(cv2.CAP_PROP_FPS))
|
|
frame_size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
|
|
|
|
if visual_difference == 0:
|
|
desired_fps = float(original_fps)
|
|
else:
|
|
desired_fps = max(float(original_fps * target_rate_of_change / visual_difference), 1.0)
|
|
|
|
desired_fps = int(desired_fps * 8) / 8
|
|
|
|
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
|
|
|
out = cv2.VideoWriter(output_filename, fourcc, desired_fps, frame_size)
|
|
|
|
while True:
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break
|
|
out.write(frame)
|
|
|
|
cap.release()
|
|
out.release()
|
|
|
|
return output_filename
|
|
|
|
|
|
def combine_videos(video_paths, base_input_directory, output_path):
|
|
import subprocess
|
|
import tempfile
|
|
|
|
# Create a temporary text file
|
|
with tempfile.NamedTemporaryFile(mode="w+t", delete=False) as temp_file:
|
|
for video_path in video_paths:
|
|
temp_file.write(f"file '{video_path}'\n")
|
|
|
|
# TODO: Check cross-platform compatibility (linux)
|
|
cmd = f"ffmpeg -f concat -safe 0 -i \"{temp_file.name}\" \"{output_path}combined{os.path.splitext(video_paths[0])[1]}\""
|
|
|
|
try:
|
|
subprocess.run(cmd, shell=True, check=True)
|
|
return output_path
|
|
except subprocess.CalledProcessError as e:
|
|
return str(e)
|
|
finally:
|
|
# Clean up the temporary file
|
|
temp_file.close()
|
|
os.remove(temp_file.name)
|
|
|
|
|
|
def transcribe_to_segs(input_filename):
|
|
import whisper as wh
|
|
import folder_paths as fp
|
|
|
|
if not os.sep in input_filename:
|
|
input_dir = fp.input_directory
|
|
full_path = os.path.join(input_dir, input_filename)
|
|
else:
|
|
full_path = input_filename
|
|
# check that the file exists
|
|
|
|
if not os.path.exists(full_path):
|
|
raise ValueError(f"File {full_path} does not exist.")
|
|
# check that it is a file
|
|
if not os.path.isfile(full_path):
|
|
raise ValueError(f"{full_path} is not a file.")
|
|
|
|
model = wh.load_model("medium.en")
|
|
|
|
res = wh.transcribe(model=model, audio=full_path, word_timestamps=True)
|
|
segs = res["segments"]
|
|
|
|
return segs
|
|
|
|
|
|
class WhisperTranscription(metaclass=ABCWidgetMetaclass):
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {"input_filename": ("STRING", {"default": "fullpath"})}}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING",)
|
|
RETURN_NAMES = ("lines_out", "words_out",)
|
|
FUNCTION = "transcribe_movie"
|
|
CATEGORY = "transcription"
|
|
|
|
def transcribe_movie(self, input_filename):
|
|
|
|
import folder_paths as fp
|
|
input_dir = fp.input_directory
|
|
full_path = os.path.join(input_dir, input_filename)
|
|
|
|
try:
|
|
segs = transcribe_to_segs(full_path)
|
|
|
|
lines_out = []
|
|
for seg in segs:
|
|
lines_out.append(f"{seg['start']:3.3f} - {seg['end']:3.3f}: {seg['text']}")
|
|
|
|
lines_out = '\n'.join(lines_out)
|
|
|
|
words_out = []
|
|
for seg in segs:
|
|
words_out.append([])
|
|
for word in seg['words']:
|
|
line = f"{word['start']:3.3f} - {word['end']:3.3f}: {word['word']}"
|
|
words_out[-1].append(line)
|
|
|
|
words_out = '\n'.join(['\n'.join(x) for x in words_out])
|
|
|
|
return (lines_out, words_out,)
|
|
|
|
except Exception as e:
|
|
print(e)
|
|
return ("Error",)
|
|
|
|
|
|
class WhisperTranscribeToFrameIDXandWords(metaclass=ABCWidgetMetaclass):
|
|
"""Just like WhisperTranscription, but returns a list of frame numbers and words instead of a string."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
# needs full path to video file
|
|
# needs fps for conversion
|
|
ret = {"required": {"input_filename": ("STRING", {"default": "fullpath"}),
|
|
"fps": ("INT", {"default": 30})}}
|
|
return ret
|
|
|
|
RETURN_TYPES = ("LIST", "LIST", "LIST",)
|
|
RETURN_NAMES = ("frame_idx", "words_out", "combined",)
|
|
FUNCTION = "transcribe_movie"
|
|
CATEGORY = "transcription"
|
|
|
|
def transcribe_movie(self, input_filename, fps):
|
|
import folder_paths as fp
|
|
input_dir = fp.input_directory
|
|
full_path = os.path.join(input_dir, input_filename)
|
|
try:
|
|
segs = transcribe_to_segs(full_path)
|
|
frame_idx = []
|
|
words_out = []
|
|
combined = []
|
|
for seg in segs:
|
|
for word in seg['words']:
|
|
start_frame = int(word['start'] * fps)
|
|
end_frame = int(word['end'] * fps)
|
|
word_text = word['word']
|
|
|
|
frame_idx.append((start_frame, end_frame))
|
|
words_out.append(word_text)
|
|
combined.append((start_frame, end_frame, word_text))
|
|
|
|
# now we fill in the gaps so the entire video is covered by some word
|
|
# we do this by finding if there is a gap then changing the frame stamps of the last word and the next word
|
|
# to cover the gap
|
|
for i in range(len(frame_idx) - 1):
|
|
if frame_idx[i][1] < frame_idx[i + 1][0]:
|
|
# there is a gap
|
|
gap_len = frame_idx[i + 1][0] - frame_idx[i][1]
|
|
to_previous_word = gap_len // 2
|
|
to_next_word = gap_len - to_previous_word
|
|
|
|
frame_idx[i] = (frame_idx[i][0], frame_idx[i][1] + to_previous_word)
|
|
frame_idx[i + 1] = (frame_idx[i + 1][0] - to_next_word, frame_idx[i + 1][1])
|
|
|
|
combined[i] = (combined[i][0], combined[i][1] + to_previous_word, combined[i][2])
|
|
combined[i + 1] = (combined[i + 1][0] - to_next_word, combined[i + 1][1], combined[i + 1][2])
|
|
|
|
# leave words_out alone
|
|
|
|
return (frame_idx, words_out, combined,)
|
|
except Exception as e:
|
|
print(e)
|
|
return ("Error",)
|
|
|
|
|
|
# @profile
|
|
def get_text_frames(word_idx, words_data, x: callable, y: callable, size: (), text_kwargs: dict, style="default"):
|
|
from .image import TextRender as text_render_class
|
|
text_render = text_render_class().render_text
|
|
|
|
start = words_data[word_idx][0]
|
|
end = words_data[word_idx][1]
|
|
word_text = words_data[word_idx][2]
|
|
frame_length = end - start
|
|
|
|
# if x and y are not callables
|
|
if not callable(x):
|
|
x = lambda i: x
|
|
if not callable(y):
|
|
y = lambda i: y
|
|
|
|
out_frames = []
|
|
for i in range(start, end):
|
|
it_dict = dict()
|
|
it_dict["text"] = word_text
|
|
it_dict["x"] = x(i / float(frame_length))
|
|
it_dict["y"] = y(i / float(frame_length))
|
|
it_dict["width"] = size[0]
|
|
it_dict["height"] = size[1]
|
|
|
|
it_dict.update(text_kwargs)
|
|
# it_dict["size"] = 128
|
|
# it_dict["color"] = "#ffffff"
|
|
# it_dict["func_only"] = False
|
|
# it_dict["stroke fill"] = "#050505"
|
|
# it_dict["stroke width"] = 5
|
|
|
|
text_img = text_render(**it_dict)
|
|
out_frames.append(text_img[0])
|
|
return out_frames
|
|
|
|
|
|
class AddTranscriptionToVideo(metaclass=ABCWidgetMetaclass):
|
|
"""
|
|
Add a transcription to a video. Requires a video and transcription data.
|
|
transcription data must be [(start_frame, end_frame, word), ...]
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
# needs full path to video file
|
|
# needs fps for conversion
|
|
ret = {"required": {"video_file": ("STRING", {"default": "fullpath"}),
|
|
"words": ("LIST", {"default": None}),
|
|
}
|
|
}
|
|
return ret
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE",)
|
|
RETURN_NAMES = ("1 fps preview", "Full Video")
|
|
FUNCTION = "add_transcription_to_video"
|
|
CATEGORY = "transcription"
|
|
|
|
# @profile
|
|
def add_transcription_to_video(self, **kwargs):
|
|
import torch
|
|
import gc
|
|
|
|
video_file = kwargs["video_file"]
|
|
words = kwargs["words"]
|
|
|
|
def create_video_frames_folder(video_file):
|
|
fnc = VideoToFramesFolderFFMPEG().handler
|
|
call_dict = dict()
|
|
call_dict["video_in"] = video_file
|
|
call_dict["folder_out"] = os.path.join(video_file, os.sep, "frames_tmp")
|
|
call_dict["format"] = "PNG"
|
|
|
|
tmp_video_folder = fnc(**call_dict)[0]
|
|
del fnc
|
|
del call_dict
|
|
gc.collect()
|
|
return tmp_video_folder
|
|
|
|
def get_video_frames(start, end):
|
|
fnc = VideoFramesFolderToImageStack().handler
|
|
call_dict = dict()
|
|
call_dict["use_subfolders"] = False
|
|
call_dict["by_date_or_name"] = "NAME"
|
|
call_dict["folder_in"] = os.path.join(video_file, os.sep, "frames_tmp")
|
|
call_dict["idx_slice_start"] = 0
|
|
call_dict["idx_slice_stop"] = end - start
|
|
call_dict["slice_idx_step"] = 1
|
|
call_dict["idx"] = start
|
|
call_dict["use_float16"] = False
|
|
|
|
ret = fnc(**call_dict)[0]
|
|
|
|
return ret
|
|
|
|
def overlay_text_on_video(video_frames, text_frames):
|
|
import torch
|
|
# Create a mask for the text frames (where text is black)
|
|
text_mask = (text_frames == 0) * 1 # Assuming black is represented as 0 in text_frames
|
|
|
|
# Use the ~ operator to invert the boolean mask
|
|
# inverted_mask = text_mask
|
|
|
|
# Overlay text_frames on video_frames using the inverted_mask
|
|
overlaid_frames = video_frames * text_mask + text_frames
|
|
|
|
return overlaid_frames
|
|
|
|
frames_out = []
|
|
num_frames = 0
|
|
frames_folder = create_video_frames_folder(video_file)
|
|
for w_idx in range(len(words)):
|
|
word_start = words[w_idx][0]
|
|
word_end = words[w_idx][1]
|
|
word = words[w_idx][2].strip()
|
|
print(f"word_start: {word_start}, word_end: {word_end}, word: {word}")
|
|
|
|
words_video_frames = get_video_frames(word_start, word_end)
|
|
i_width = words_video_frames.shape[2]
|
|
i_height = words_video_frames.shape[1]
|
|
text_frames = get_text_frames(word_idx=w_idx,
|
|
words_data=words,
|
|
x=lambda i: 0,
|
|
y=lambda i: 0,
|
|
size=(i_width, i_height),
|
|
text_kwargs={"size": 128, "color": "#ffffff", "func_only": False,
|
|
"stroke fill": "#050505", "stroke width": 5,
|
|
"font_name": "c:/Windows/Fonts/arial.ttf"})
|
|
|
|
text_frames = torch.cat(text_frames)
|
|
|
|
# Overlay text on video frames
|
|
video_frames_with_text = overlay_text_on_video(words_video_frames, text_frames)
|
|
frames_out.append(video_frames_with_text)
|
|
num_frames += video_frames_with_text.shape[0]
|
|
gc.collect()
|
|
|
|
# free memory
|
|
del words_video_frames
|
|
del text_frames
|
|
del video_frames_with_text
|
|
|
|
num_stks = len(frames_out)
|
|
# write to disk
|
|
for i, stk in enumerate(frames_out):
|
|
torch.save(stk, f"tmp_str_{i}.pt")
|
|
del frames_out
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
# pre allocate tensor
|
|
y = torch.zeros((num_frames, i_height, i_width, 3), dtype=torch.float32, device="cpu")
|
|
# load each save and put into tensor
|
|
new_start = 0
|
|
for i in range(num_stks):
|
|
stk = torch.load(f"tmp_str_{i}.pt")
|
|
y[new_start:new_start + stk.shape[0]] = stk
|
|
new_start = new_start + stk.shape[0]
|
|
del stk
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
preview = y[::30].clone()
|
|
|
|
return (preview, y,)
|
|
|
|
|
|
class TransformImageStack(metaclass=ABCWidgetMetaclass):
|
|
"""
|
|
Transform a stack of images
|
|
given x change (in pixels)
|
|
given y change (in pixels)
|
|
|
|
|
|
"""
|
|
|
|
# TODO: support sheer,rotate,scale
|
|
# TODO: support transparancy
|
|
# TODO: support fill color specification
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
ret = {"required": {"ImageStack": ("IMAGE", {"default": None}),
|
|
"pan x": ("INT", {"default": 0, "min": -2048}),
|
|
"pan y": ("INT", {"default": 0, "min": -2048}),
|
|
},
|
|
}
|
|
return ret
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "transform_image_stack"
|
|
|
|
def transform_image_stack(self, **kwargs):
|
|
import os
|
|
import torch
|
|
from PIL import Image
|
|
import torchvision.transforms as transforms
|
|
import torchvision.transforms.functional as TF
|
|
|
|
x_image = kwargs["ImageStack"]
|
|
x_image = x_image.detach().clone()
|
|
x = kwargs["pan x"]
|
|
y = kwargs["pan y"]
|
|
|
|
time_index, H, W, C = x_image.shape # Assuming x_image has shape (time_index, H, W, C)
|
|
out_list = []
|
|
|
|
stop_x = x
|
|
stop_y = y
|
|
dx_step = stop_x / time_index - 1
|
|
dy_step = stop_y / time_index - 1
|
|
|
|
for t in range(time_index):
|
|
tx = dx_step * t
|
|
ty = dy_step * t
|
|
|
|
current_frame = x_image[t].detach().clone()
|
|
# Change from (H, W, C) to (C, H, W)
|
|
current_frame = current_frame.permute(2, 0, 1)
|
|
|
|
transformed_frame = TF.affine(current_frame, angle=0, translate=(tx, ty), scale=1,
|
|
shear=[0, 0]).detach().clone()
|
|
|
|
# Convert back to (H, W, C)
|
|
transformed_frame = transformed_frame.permute(1, 2, 0)
|
|
|
|
# Store in the list
|
|
x_image[t] = transformed_frame
|
|
|
|
return (x_image.detach().clone(),)
|
|
|
|
|
|
class PanByStep(metaclass=ABCWidgetMetaclass):
|
|
"""
|
|
Transform a stack of images
|
|
given x change (in pixels)
|
|
given y change (in pixels)
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
ret = {"required": {"ImageStack": ("IMAGE", {"default": None}),
|
|
"pan x": ("INT", {"default": 0, "min": -2048}),
|
|
"pan y": ("INT", {"default": 0, "min": -2048}),
|
|
"wrap": ("BOOLEAN", {"default": False}),
|
|
},
|
|
}
|
|
return ret
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "pan_by_step"
|
|
|
|
def pan_by_step(self, **kwargs):
|
|
import os
|
|
import torch
|
|
from PIL import Image
|
|
import torchvision.transforms as transforms
|
|
import torchvision.transforms.functional as TF
|
|
|
|
x_image = kwargs["ImageStack"]
|
|
x_image = x_image.detach().clone()
|
|
x = kwargs["pan x"]
|
|
y = kwargs["pan y"]
|
|
|
|
wrap = kwargs["wrap"]
|
|
|
|
time_index, H, W, C = x_image.shape # Assuming x_image has shape (time_index, H, W, C)
|
|
out_list = []
|
|
|
|
stop_x = x
|
|
stop_y = y
|
|
dx_step = 0 if x == 0 else stop_x / time_index - 1
|
|
dy_step = stop_y / time_index - 1
|
|
|
|
for t in range(time_index):
|
|
tx = dx_step * t
|
|
ty = dy_step * t
|
|
|
|
current_frame = x_image[t].detach().clone()
|
|
# Change from (H, W, C) to (C, H, W)
|
|
current_frame = current_frame.permute(2, 0, 1)
|
|
|
|
# transformed_frame = TF.affine(current_frame, angle=0, translate=(tx, ty), scale=1,
|
|
# shear=[0, 0]).detach().clone()
|
|
|
|
# If wrap is True, apply the wrapping operation
|
|
if wrap:
|
|
transformed_frame = current_frame.roll(
|
|
(int(tx) % W, int(ty) % H), dims=(2, 1)
|
|
)
|
|
else:
|
|
transformed_frame = TF.affine(current_frame, angle=0, translate=(tx, ty), scale=1,
|
|
shear=[0, 0]).detach().clone()
|
|
|
|
# Convert back to (H, W, C)
|
|
transformed_frame = transformed_frame.permute(1, 2, 0)
|
|
|
|
# Store in the list
|
|
x_image[t] = transformed_frame
|
|
|
|
return (x_image.detach().clone(),)
|
|
|
|
|
|
class VideoFolderProvider(metaclass=ABCWidgetMetaclass):
|
|
"""
|
|
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": None}),
|
|
},
|
|
"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 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")
|
|
|
|
# if the folder does not exist, create the folder
|
|
|
|
return (kwargs["folder_path"],)
|
|
|
|
|
|
class AudioFolderProvider(metaclass=ABCWidgetMetaclass):
|
|
"""
|
|
given a folder path, creates the folder or empties it entirely returns the folder path as
|
|
a string but as the node type "AUDIOFOLDER"
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
ret = {"required": {"folder_path": ("STRING", {"default": None}),
|
|
},
|
|
}
|
|
return ret
|
|
|
|
RETURN_TYPES = ("AUDIOFOLDER",)
|
|
FUNCTION = "audio_folder_provider"
|
|
|
|
def audio_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 audio files is not accessible")
|
|
|
|
# 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)
|
|
|
|
# if the folder does not exist, create the folder
|
|
|
|
return (kwargs["folder_path"],)
|
|
|
|
|
|
class VideoDefinitionProvider(metaclass=ABCWidgetMetaclass):
|
|
"""
|
|
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
|
|
ret = adict()
|
|
ret["video_folder"] = kwargs["video_folder"]
|
|
ret["height"] = kwargs["height"]
|
|
ret["width"] = kwargs["width"]
|
|
ret["fps"] = kwargs["fps"]
|
|
ret["frame format"] = kwargs["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 dont 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,)
|
|
|
|
|
|
class AudioDefinitionProvider(metaclass=ABCWidgetMetaclass):
|
|
"""
|
|
given a AUDIOFOLDER node, and inputs for sample rate, and channels, creates a AUDIOFOLDERDEF variable
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
ret = {"required": {"audio_folder": ("AUDIOFOLDER", {"default": None}),
|
|
"sample rate": ("INT", {"default": 44100}),
|
|
"channels": ("INT", {"default": 2}),
|
|
"audio format": (["WAV", "MP3"], {"default": "WAV"}),
|
|
},
|
|
}
|
|
return ret
|
|
|
|
RETURN_TYPES = ("AUDIO_FOLDER_DEF",)
|
|
FUNCTION = "audio_definition_provider"
|
|
|
|
def get_num_audio_files_that_match(self, audio_folder):
|
|
audio_files = self.get_all_audio_files(audio_folder)
|
|
return len(audio_files)
|
|
|
|
def get_all_audio_files(self, audio_folder):
|
|
import os
|
|
audio_files = os.listdir(audio_folder)
|
|
audio_files = [os.path.join(audio_folder, x) for x in audio_files if "audio_" in x]
|
|
return audio_files
|
|
|
|
def audio_definition_provider(self, **kwargs):
|
|
from addict import Dict as adict
|
|
ret = adict()
|
|
ret["audio_folder"] = kwargs["audio_folder"]
|
|
ret["sample rate"] = kwargs["sample rate"]
|
|
ret["channels"] = kwargs["channels"]
|
|
ret["audio format"] = kwargs["audio format"]
|
|
# a function that gets the frame count from the disk
|
|
ret.get_file_count = lambda: self.get_num_audio_files_that_match(ret["audio_folder"])
|
|
# a function that gets the next file name and full path dont forget to zfill
|
|
ret.get_next_file_name = lambda: f"{ret['audio_folder']}{os.sep}audio_{str(ret.get_file_count()).zfill(10)}.{ret['audio format'].lower()}"
|
|
ret.get_last_file_name = lambda: f"{ret['audio_folder']}{os.sep}audio_{str(ret.get_file_count() - 1).zfill(10)}.{ret['audio format'].lower()}"
|
|
ret.get_all_audio_files = lambda: self.get_all_audio_files(ret["audio_folder"])
|
|
|
|
return (ret,)
|
|
|
|
|
|
class AddAudioToAudioFolder(metaclass=ABCWidgetMetaclass):
|
|
"""
|
|
given a AUDIO_FOLDER_DEF node, and a list of audio files, adds the audio files to the audio folder
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
ret = {"required": {"audio_folder_def": ("AUDIO_FOLDER_DEF", {"default": None}),
|
|
"AUDIO": ("LIST", {"default": None}),
|
|
"delete after": ("BOOLEAN", {"default": False}),
|
|
},
|
|
}
|
|
return ret
|
|
|
|
RETURN_TYPES = ("AUDIO_FOLDER_DEF",)
|
|
FUNCTION = "add_audio_to_audio_folder"
|
|
|
|
def add_audio_to_audio_folder(self, **kwargs):
|
|
# check each file in the list to make sure we can convert it to the correct format
|
|
import shutil
|
|
import os
|
|
|
|
ffmpeg_exe = ffmpeg.get_ffmpeg_exe()
|
|
|
|
audio_folder_def = kwargs["audio_folder_def"]
|
|
audio_files = kwargs["AUDIO"]
|
|
|
|
delete = kwargs["delete after"]
|
|
|
|
for audio_file in audio_files:
|
|
# use ffmpeg_exe to convert the audio file to the correct format which is audio_folder_def["audio format"]
|
|
# use the audio_folder_def["audio format"] to determine the output file name
|
|
# use the audio_folder_def["audio folder"] to determine the output folder
|
|
# use the audio_folder_def["sample rate"] to determine the sample rate
|
|
# use the audio_folder_def["channels"] to determine the number of channels
|
|
|
|
# get the output file name
|
|
output_file_name = audio_folder_def.get_next_file_name()
|
|
input_file_name = audio_file
|
|
|
|
CMD = f"{ffmpeg_exe} -y -i {input_file_name} -ar {audio_folder_def['sample rate']} -ac {audio_folder_def['channels']} {output_file_name}"
|
|
|
|
os.system(CMD)
|
|
if delete:
|
|
os.remove(input_file_name)
|
|
|
|
return (audio_folder_def,)
|
|
|
|
|
|
class AddSilenceToAudioFolder(metaclass=ABCWidgetMetaclass):
|
|
"""create and add a new audio file that is silence for the specified duration"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
ret = {"required": {"audio_folder_def": ("AUDIO_FOLDER_DEF", {"default": None}),
|
|
"duration": ("FLOAT", {"default": 1.0}),
|
|
},
|
|
}
|
|
return ret
|
|
|
|
RETURN_TYPES = ("AUDIO_FOLDER_DEF",)
|
|
FUNCTION = "add_silence_to_audio_folder"
|
|
|
|
def add_silence_to_audio_folder(self, **kwargs):
|
|
# this is to be multi-platform
|
|
import shutil
|
|
import os
|
|
|
|
audio_folder_def = kwargs["audio_folder_def"]
|
|
duration = kwargs["duration"]
|
|
|
|
# get the output file name
|
|
output_file_name = audio_folder_def.get_next_file_name()
|
|
sample_rate = audio_folder_def["sample rate"]
|
|
channels = audio_folder_def["channels"]
|
|
fmt = audio_folder_def["audio format"]
|
|
|
|
# create the silence audio file
|
|
|
|
if fmt == "WAV":
|
|
fmt_part = "pcm_s16le"
|
|
elif fmt == "MP3":
|
|
fmt_part = "mp3"
|
|
else:
|
|
fmt_part = "pcm_s16le"
|
|
|
|
ffmpeg_exe = ffmpeg.get_ffmpeg_exe()
|
|
# "-i anullsrc" is the source of the silence this will work on all platforms
|
|
# "-c:a" is the codec for the audio file
|
|
# "-f" is the format of the audio file, "lavfi" is equivalent to "ffmpeg filter"
|
|
# "-i anullsrc=r=44100:cl=2:d=1" is the source of the silence
|
|
|
|
# CMD = f"{ffmpeg_exe} -f lavfi -i aevalsrc=0:d={duration}:cl={channels}:s={sample_rate} -c:a {fmt_part} {output_file_name}"
|
|
# CMD = f"{ffmpeg_exe} -f lavfi -i anullsrc=channel_layout={channels}:sample_rate={sample_rate}:duration={duration} -c:a {fmt_part} {output_file_name}"
|
|
CMD = f"{ffmpeg_exe} -f lavfi -i aevalsrc=0:sample_rate={sample_rate}:channel_layout={channels},atrim=0:{duration} -c:a {fmt_part} {output_file_name}"
|
|
# now we call it in a way to that we will know if there is an error
|
|
|
|
import subprocess
|
|
# the plan is to call the command and wait for it to finish
|
|
process = subprocess.Popen(CMD, shell=True, stderr=subprocess.PIPE)
|
|
process.wait()
|
|
# check if the process was successful
|
|
if process.returncode != 0:
|
|
# if the process was not successful raise an exception
|
|
raise Exception(f"ffmpeg failed with error code {process.returncode}")
|
|
|
|
return (audio_folder_def,)
|
|
|
|
|
|
class GetAudioFilesLengthsFromAudioFolder(metaclass=ABCWidgetMetaclass):
|
|
"""Get the lengths of all audio files in a given audio folder"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {"audio_folder_def": ("AUDIO_FOLDER_DEF", {"default": None})}}
|
|
|
|
RETURN_TYPES = ("AUDIO_FOLDER_DEF", "LIST",)
|
|
RETURN_NAMES = ("audio_folder_def", "[lens(seconds)]",)
|
|
FUNCTION = "get_audio_files_lengths"
|
|
|
|
def get_audio_files_lengths(self, audio_folder_def):
|
|
import os
|
|
import torchaudio
|
|
|
|
audio_folder = audio_folder_def["audio_folder"]
|
|
audio_files = os.listdir(audio_folder)
|
|
audio_files_lengths = []
|
|
|
|
for audio_file in audio_files:
|
|
audio_path = os.path.join(audio_folder, audio_file)
|
|
waveform, sample_rate = torchaudio.load(audio_path)
|
|
duration = waveform.shape[1] / sample_rate
|
|
audio_files_lengths.append(duration)
|
|
|
|
return (audio_folder_def, audio_files_lengths,)
|
|
|
|
|
|
class AudioFolderRenderAudio(metaclass=ABCWidgetMetaclass):
|
|
""" Joins all the audio files """
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"audio_folder_def": ("AUDIO_FOLDER_DEF", {"default": None}),
|
|
"output_path": ("STRING", {"multiline": False, "default": ""}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("output_path",)
|
|
FUNCTION = 'join_audio_files'
|
|
CATEGORY = 'ETK/audio'
|
|
|
|
def join_audio_files(self, **kwargs):
|
|
import librosa
|
|
import copy
|
|
import soundfile as sf
|
|
|
|
audio_folder_def = kwargs['audio_folder_def']
|
|
kwargs = copy.deepcopy(kwargs)
|
|
paths = audio_folder_def.get_all_audio_files()
|
|
output_path = kwargs['output_path']
|
|
|
|
# Read the first file to get metadata
|
|
audio_data_list = []
|
|
for path in paths:
|
|
audio_data, sr = librosa.load(path, sr=None)
|
|
audio_data_list.append(audio_data)
|
|
|
|
# Concatenate all audio data
|
|
joined_audio_data = np.concatenate(audio_data_list)
|
|
|
|
# TODO: need to implement writing in the correct format
|
|
# Save the concatenated audio
|
|
sf.write(output_path, joined_audio_data, sr)
|
|
|
|
return (output_path,)
|
|
|
|
|
|
import subprocess
|
|
import platform
|
|
|
|
import subprocess
|
|
import platform
|
|
import os
|
|
|
|
|
|
def duplicate_file(source_file, target_files):
|
|
"""
|
|
Duplicates a given source file into multiple target files using a single CLI command
|
|
in a non-blocking way. The base directory is automatically determined from the
|
|
source file's location.
|
|
|
|
Args:
|
|
- source_file (str): Path to the source file to be duplicated.
|
|
- target_files (list of str): A list of target file paths where the source file should be duplicated.
|
|
|
|
Returns:
|
|
- None
|
|
"""
|
|
# Determine the base directory from the source file's directory
|
|
base_dir = os.path.dirname(os.path.abspath(source_file))
|
|
|
|
# Change the current working directory to the base directory
|
|
os.chdir(base_dir)
|
|
|
|
# Get the source file name (relative to the base directory)
|
|
source_file_name = os.path.basename(source_file)
|
|
|
|
# get the relative paths of the target files
|
|
target_files = [os.path.relpath(target_file, base_dir) for target_file in target_files]
|
|
|
|
# Determine the shell command based on the operating system
|
|
copy_command = 'copy' if os.name == 'nt' else 'cp'
|
|
|
|
# Construct the command string
|
|
command_parts = [f'{copy_command} "{source_file_name}" "{target}"' for target in target_files]
|
|
command = ' && '.join(command_parts)
|
|
|
|
command = command.replace("\"", "")
|
|
command = command.replace("\'", "")
|
|
# print(len(command))
|
|
|
|
try:
|
|
# Execute the command in a non-blocking way
|
|
process = subprocess.Popen(command, shell=True, stderr=subprocess.DEVNULL)
|
|
process.wait()
|
|
|
|
except subprocess.CalledProcessError as e:
|
|
pass
|
|
except ValueError as e:
|
|
pass
|
|
except Exception as e:
|
|
pass
|
|
|
|
if "process" not in locals():
|
|
pass
|
|
if "process" not in locals() or process.returncode != 0:
|
|
# If the command is too long, split it into multiple commands
|
|
# split the command into commands that do 64 at a time
|
|
command_parts = [command_parts[i:i + 64] for i in range(0, len(command_parts), 64)]
|
|
processes = []
|
|
for command_part in command_parts:
|
|
command = ' && '.join(command_part)
|
|
process = subprocess.Popen(command, shell=True, stderr=subprocess.DEVNULL)
|
|
processes.append(process)
|
|
|
|
for process in processes:
|
|
process.wait()
|
|
|
|
return process
|
|
|
|
|
|
class AddFramesToVideoFolder(metaclass=ABCWidgetMetaclass):
|
|
"""
|
|
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):
|
|
from torchvision.utils import save_image
|
|
import shutil
|
|
from PIL import Image
|
|
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
|
|
|
|
if current_frame.dtype != torch.uint8: # handle float type frames
|
|
ndarr = current_frame.permute(2, 0, 1).mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu",
|
|
torch.uint8).numpy()
|
|
else: # handle uint8 frames
|
|
ndarr = current_frame.permute(2, 0, 1).permute(1, 2, 0).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
|
|
num_frames = frames.shape[0]
|
|
previous_frame = None
|
|
frame_fp = None
|
|
frames_area_all_the_same = torch.all(frames[0] == frames)
|
|
if frames_area_all_the_same and frames.shape[0] > 1:
|
|
# get a list of file names using reservenexframe
|
|
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 glob import glob
|
|
from PIL import Image
|
|
import os
|
|
import torch
|
|
import gc
|
|
from torchvision.transforms.functional import resize
|
|
import psutil
|
|
|
|
initial_memory = psutil.virtual_memory().available
|
|
print(f"Initial free memory: {initial_memory / (1024 * 1024):.2f} MB")
|
|
|
|
import shutil
|
|
|
|
# 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 = lambda x: resize(x, (height, width), interpolation=rescale_type)
|
|
frames = rescale_method(frames)
|
|
else:
|
|
rescale_method = lambda x: lanczos_resize(x, (width, height))
|
|
frames = rescale_method(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, 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()
|
|
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"],)
|
|
|
|
|
|
class RepeatFrames(metaclass=ABCWidgetMetaclass):
|
|
"""
|
|
does nothing more than repeat the frames a number of times
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
ret = {"required": {"IMAGE": ("IMAGE", {"default": None}),
|
|
"repeat": ("INT", {"default": 1}),
|
|
},
|
|
}
|
|
return ret
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("repeated images",)
|
|
FUNCTION = "repeat_frames"
|
|
|
|
def repeat_frames(self, **kwargs):
|
|
import torch
|
|
frames = kwargs["IMAGE"]
|
|
repeat = kwargs["repeat"]
|
|
return (frames.repeat(repeat, 1, 1, 1),)
|
|
|
|
|
|
class HoldFrameForSecsInVideoFolder(metaclass=ABCWidgetMetaclass):
|
|
"""
|
|
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):
|
|
import torch
|
|
import asyncio
|
|
|
|
vide_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 = vide_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"],)
|
|
|
|
|
|
class Listed_HoldFramesForSecsInVideoFolder(metaclass=ABCWidgetMetaclass):
|
|
"""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
|
|
|
|
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"],)
|
|
|
|
|
|
class VideoFolderDefInfo(metaclass=ABCWidgetMetaclass):
|
|
"""
|
|
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,)
|
|
|
|
|
|
class FadeToInVideoFolder(metaclass=ABCWidgetMetaclass):
|
|
"""
|
|
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
|
|
|
|
# usee 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, H, W, C)
|
|
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
|
|
new_frame = (1 - alpha) * base_frame + alpha * new_frame
|
|
new_frame = new_frame.to("cpu")
|
|
transition_stack[i] = new_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"],)
|
|
|
|
|
|
class ListedSVDConditioning(metaclass=ABCWidgetMetaclass):
|
|
"""
|
|
repeats SVD_img2vid_Conditioning for each set of params
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
base = SVD_img2vid_Conditioning.INPUT_TYPES()
|
|
base["required"]["init_image"] = ("LIST", {"default": None})
|
|
# base["required"]["clip_vision"] = ("LIST", {"default": None})
|
|
base["required"]["video_frames"] = ("LIST", {"default": None})
|
|
|
|
return base
|
|
|
|
RETURN_TYPES = ("LIST", "LIST", "LIST")
|
|
RETURN_NAMES = ("positive", "negative", "latent")
|
|
|
|
FUNCTION = "listed_SVD_conditioning"
|
|
|
|
def listed_SVD_conditioning(self, **kwargs):
|
|
# just like Listed_HoldFramesForSecsInVideoFolder, we will call SVD_img2vid_Conditioning for each set of params
|
|
|
|
positives = []
|
|
negatives = []
|
|
latents = []
|
|
|
|
init_images = kwargs["init_image"]
|
|
clip_vision = kwargs["clip_vision"]
|
|
# vae
|
|
vae = kwargs["vae"]
|
|
# width and height
|
|
width = kwargs["width"]
|
|
height = kwargs["height"]
|
|
video_frames = kwargs["video_frames"]
|
|
motion_bucket_id = kwargs["motion_bucket_id"]
|
|
fps = kwargs["fps"]
|
|
augmentation_level = kwargs["augmentation_level"]
|
|
|
|
lens = [len(init_images), len(video_frames)]
|
|
|
|
# check if one length is divisible by the other
|
|
if lens[0] % lens[1] != 0 and lens[1] % lens[0] != 0:
|
|
raise ValueError("The length of init_image and video_frames must be divisible by each other")
|
|
|
|
# repeat the shorter
|
|
if lens[0] < lens[1]:
|
|
init_images = init_images * (lens[1] // lens[0])
|
|
else:
|
|
video_frames = video_frames * (lens[0] // lens[1])
|
|
|
|
# at this point the lengths should be the same
|
|
|
|
for init_image, video_frame in zip(init_images, video_frames):
|
|
pos, neg, latent = SVD_img2vid_Conditioning().encode(
|
|
clip_vision=clip_vision,
|
|
init_image=init_image,
|
|
video_frames=video_frame,
|
|
vae=vae,
|
|
width=width,
|
|
height=height,
|
|
motion_bucket_id=motion_bucket_id,
|
|
fps=fps,
|
|
augmentation_level=augmentation_level
|
|
)
|
|
|
|
positives.append(pos)
|
|
negatives.append(neg)
|
|
latents.append(latent)
|
|
|
|
return (positives, negatives, latents,)
|
|
|
|
|
|
class VideoFolderFFMPEGVideoInterpolate(metaclass=ABCWidgetMetaclass):
|
|
"""
|
|
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)
|
|
|
|
# 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)
|
|
|
|
# put 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.reserve_next_frame() for _ in range(target_n_frames)]
|
|
output_frame_names = [video_folder_def.get_next_file_name(i) for i in range(target_n_frames)]
|
|
|
|
# delete the frames on the list
|
|
#for frame_name in output_frame_names:
|
|
# os.remove(frame_name)
|
|
|
|
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",
|
|
"-i", input_frame_list_file,
|
|
"-vsync", "0", "-filter_complex",
|
|
# f"[0:v]minterpolate=fps={target_fps}:mi_mode=mci:mc_mode=aobmc:me_mode=bidir:me=ds:mb_size=16:search_param=8:vsbmc=1[v]",
|
|
# f"[0:v]framerate=fps={target_fps}:interp_start=0:interp_end=255:scene=0[v]",
|
|
# f"[0:v]minterpolate=fps={target_fps}:vsbmc=1[v]",
|
|
f"[0:v]minterpolate=fps={target_fps}:mi_mode=blend[v]",
|
|
"-map", "[v]", "-q:v", "2", "-r", str(target_fps), "-frames", 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,)
|