nodes/TashaSkyUp-EternalKernelPytorchNodes
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

added whisper-openai to requirements as well as moviepy

Browse Source 
started creating whisper transcription node

modified how imports are handled to avoid importing the same moduels twice under different name spaces.

added basic.py to house simple functions that work on primitives.

NOTE: this update relies on change to /web/extensions that are not included in this commit, a future commit will include an install process for these changes.
This commit is contained in:
Tasha Upchurch
2023-09-07 22:02:39 -04:00
parent 7389631b65
commit aed1618bf7
4 changed files with 591 additions and 139 deletions
Download Patch File Download Diff File Expand all files Collapse all files
basic.py
+179
View File
@@ -0,0 +1,179 @@
class MultilineStringNode:
"""Noe that displays multiline text box"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"Text": ("STRING", {
"default": "",
"multiline": True,
}),
}
}
RETURN_TYPES = ("STRING", "INT", "FLOAT", "LIST")
FUNCTION = "get_value"
CATEGORY = "text"
def get_value(self, Text):
return (Text,)
class StrBreakout:
""" node that breaks out a string into a:
- Boolean
- int
- float
- [string]
- splitlines()
- split(",")
the first value is a description of the transformations tried
"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"Text": ("STRING", {
"default": "",
}),
}
}
RETURN_TYPES = ("STRING", "BOOL", "INT", "FLOAT", "LIST", "LIST", "LIST",)
FUNCTION = "get_value"
CATEGORY = "text"
def get_value(self, Text):
out_bool = False
out_int = 0
out_float = 0.0
out_list = []
out_lines = []
out_comma = []
notes = []
try:
out_bool = bool(Text)
notes.append(f"bool worked")
except Exception as e:
notes.append(f"bool failed: {e}")
try:
out_int = int(Text)
notes.append(f"int worked")
except Exception as e:
notes.append(f"int failed: {e}")
pass
try:
out_float = float(Text)
notes.append(f"float worked")
except Exception as e:
pass
try:
out_list = list(Text)
notes.append(f"list worked")
except Exception as e:
notes.append(f"list failed: {e}")
pass
try:
out_lines = Text.splitlines()
notes.append(f"splitlines worked")
except Exception as e:
notes.append(f"splitlines failed: {e}")
pass
try:
out_comma = Text.split(",")
notes.append(f"split(',') worked")
except Exception as e:
notes.append(f"split(',') failed: {e}")
pass
return ("\n".join(notes), out_bool, out_int, out_float, out_list, out_lines, out_comma,)
class IntBreakout:
""" node that breaks out an int into a:
- Boolean
- str(int)
- float
- [int]
- list(range(int))
- [None] * int
the first value is a description of the transformations tried
"""
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"Int": ("INT", {
"default": 0,
}),
}
}
RETURN_TYPES = ("STRING", "BOOL", "STRING", "FLOAT", "LIST", "LIST", "LIST",)
FUNCTION = "get_value"
CATEGORY = "text"
def get_value(self, Int):
out_bool = False
out_str = ""
out_float = 0.0
out_list = []
out_range = []
out_none = []
notes = []
try:
out_bool = bool(Int)
notes.append(f"bool worked")
except Exception as e:
notes.append(f"bool failed: {e}")
try:
out_str = str(Int)
notes.append(f"str worked")
except Exception as e:
notes.append(f"str failed: {e}")
pass
try:
out_float = float(Int)
notes.append(f"float worked")
except Exception as e:
pass
try:
out_list = list(Int)
notes.append(f"list worked")
except Exception as e:
notes.append(f"list failed: {e}")
pass
try:
out_range = list(range(Int))
notes.append(f"range worked")
except Exception as e:
notes.append(f"range failed: {e}")
pass
try:
out_none = [None] * Int
notes.append(f"[None] * worked")
except Exception as e:
notes.append(f"[None] * failed: {e}")
pass
return ("\n".join(notes), out_bool, out_str, out_float, out_list, out_range, out_none,)
NODE_CLASS_MAPPINGS = {}
NODE_CLASS_MAPPINGS["MultilineStringNode"] = MultilineStringNode
NODE_CLASS_MAPPINGS["StrBreakout"] = StrBreakout
NODE_CLASS_MAPPINGS["IntBreakout"] = IntBreakout
image.py
+107 -50
View File
@@ -1,5 +1,11 @@
import os
try:
from line_profiler_pycharm import profile
except ImportError as e:
print("ETK> line_profiler_pycharm not found, skipping it")
profile = lambda x: x
try:
from utils import etk_deep_copy as deepcopy
except ImportError as e:
@@ -22,8 +28,8 @@ if testing:
pass
else:
try:
from nodes import CLIPTextEncode, VAEEncode, VAEDecode, KSampler, CheckpointLoaderSimple, EmptyLatentImage, \
SaveImage
from nodes import VAEDecode, KSampler, CheckpointLoaderSimple, EmptyLatentImage
from nodes_examples import CLIPTextEncode, VAEEncode, SaveImage
from nodes import common_ksampler
@@ -213,11 +219,17 @@ class TinyTxtToImg:
RETURN_NAMES = ("image", "FUNC", "latent", "positive", "negative",)
FUNCTION = "tinytxt2img"
@profile
def tinytxt2img(self, **kwargs):
""" use the imports from nodes to generate an image from text """
import random
import json
import gc
# Create a stream for computation
gpu_stream = torch.cuda.Stream(device="cuda")
cpu_stream = torch.Stream()
# kwargs = deepcopy(kwargs)
@@ -230,6 +242,8 @@ class TinyTxtToImg:
overrides = kwargs.get("overrides", "")
name = kwargs.get("name", "tinytxt2img")
FUNC = kwargs.get("FUNC", None)
render_what = kwargs.get("render_what", "image")
self.latent_image = kwargs.get("latent", None)
self.pos_cond = kwargs.get("pos_cond", None)
self.neg_cond = kwargs.get("neg_cond", None)
@@ -274,7 +288,6 @@ class TinyTxtToImg:
TinyTxtToImg.share_vae
)
self.cfg = 8
self.sampler_name = comfy.samplers.KSampler.SAMPLERS[0]
self.scheduler = comfy.samplers.KSampler.SCHEDULERS[0]
@@ -321,27 +334,54 @@ class TinyTxtToImg:
print(e)
raise ValueError("advanced clip encoder failed")
#if not self.latent_image == None:
#self.latent_image = EmptyLatentImage.generate(None, self.width, self.height, self.batch_size)[0]
if self.latent_image == None:
self.latent_image = EmptyLatentImage.generate(None, self.width, self.height, self.batch_size)[0]
if "latent" in render_what or "all" in render_what:
self.KSampler = KSampler()
self.samples = self.KSampler.sample(
model=self.mdl,
seed=self.seed,
steps=self.steps,
cfg=self.cfg,
sampler_name=self.sampler_name,
scheduler=self.scheduler,
latent_image=self.latent_image,
pos_cond=self.pos_cond,
neg_cond=self.neg_cond,
positive=self.pos_cond,
negative=self.neg_cond,
denoise=self.denoise,
one_seed_per_batch=self.one_seed_per_batch)[0]
if "image" in render_what or "all" in render_what:
# Move the samples to the GPU
self.samples["samples"] = self.samples["samples"].to("cuda", non_blocking=True)
# Move VAE model to the GPU asynchronously
self.vae.first_stage_model.to("cuda:0", non_blocking=True)
self.vae.device = torch.device(torch.device("cuda:0"))
torch.cuda.empty_cache()
decoded_images = self.vae.decode(self.samples["samples"])
image = decoded_images.to("cpu", non_blocking=True)
self.vae.first_stage_model.to("cpu", non_blocking=True)
self.vae.device = torch.device(torch.device("cpu"))
torch.cuda.empty_cache()
gc.collect()
# Perform VAEDecode using samples in the computation stream
# with torch.cpu.stream(cpu_stream):
# decoded_images = self.vae.decode(self.samples["samples"])
# Synchronize computation stream before moving model back to CPU
# torch.cuda.synchronize()
# Move VAE model back to CPU asynchronously
# self.vae.first_stage_model.to("cpu", non_blocking=True)
# Convert the decoded image to CPU
self.KSampler = KSampler()
self.samples = self.KSampler.sample(
model=self.mdl,
seed=self.seed,
steps=self.steps,
cfg=self.cfg,
sampler_name=self.sampler_name,
scheduler=self.scheduler,
latent_image=self.latent_image,
pos_cond=self.pos_cond,
neg_cond=self.neg_cond,
positive=self.pos_cond,
negative=self.neg_cond,
denoise=self.denoise,
one_seed_per_batch=self.one_seed_per_batch)[0]
# self.vae.offload_device = torch.device("cuda")
image = VAEDecode.decode(None, self.vae, self.samples)[0]
image = image.detach().cpu()
self.FUNC = self.tinytxt2img
self.ARGS = kwargs
ret = (image,
@@ -1268,7 +1308,8 @@ class StackImages:
"image2": ("IMAGE",),
"number_of_images": ("INT", {"default": 2, "min": 2, "max": 100}),
# an optionbox with "all" or "upto", model after "one seed per batch" in the sampler
"mode": (["all", "upto"], {"default": "all"}),
"mode": (["all", "one_for_one", "upto"], {"default": "all"}),
"number_of_repeats": ("INT", {"default": 1, "min": 1, "max": 100}),
}, }
@@ -1277,14 +1318,12 @@ class StackImages:
CATEGORY = "ETK/image"
def stackme(self, image1, image2=None, number_of_images=2, mode="all"):
def stackme(self, image1, image2=None, number_of_images=2, mode="all", number_of_repeats=1):
# import modules to help clear memory
import gc
import torch
import torch.cuda
"""
stacks two images on top of each other
images are (B, W, H, C) torch tensors
@@ -1295,36 +1334,51 @@ class StackImages:
if image2 is not None:
image2 = image2.clone()
#gc.collect()
#torch.cuda.empty_cache()
# gc.collect()
# torch.cuda.empty_cache()
del gc.garbage[:]
gc.collect()
torch.cuda.empty_cache()
# apply number of repeats
if image2 is not None:
if mode == "all" and image2 is not None:
image2 = image2.clone()
stacked = torch.cat((image1, image2), dim=0)
else:
stacked = torch.cat((image1, image1), dim=0)
out = torch.cat((image1, image2), dim=0)
elif mode == "all" and image2 is None:
out = image1.clone()
elif mode == "one_for_one" and image2 is not None:
# for each image in image1, add each to the output
out_size = min(image1.shape[0], image2.shape[0])
holder = torch.zeros((out_size, image1.shape[1], image1.shape[2], image1.shape[3]))
for i in range(out_size):
holder[i, :, :, :] = image2[i, :, :, :]
out = holder
elif mode == "one_for_one" and image2 is None:
# this might mean that the user whats each index repeated num times
out_size = number_of_repeats * image1.shape[0]
out = torch.zeros((out_size, image1.shape[1], image1.shape[2], image1.shape[3]))
for i in range(out_size):
out[i, :, :, :] = image1[i % image1.shape[0], :, :, :]
if mode == "all":
if number_of_images > 2:
while stacked.shape[0] < number_of_images:
stacked = torch.cat((stacked, stacked), dim=0)
print(stacked.shape)
out = stacked[:number_of_images, :, :, :]
else:
# stack them all together regardless of the first dimension
out = torch.cat((image1, image2), dim=0)
print(stacked.shape)
elif mode == "upto" and image2 is not None:
# stack them all together but then only return the first number_of_images
out = torch.cat((image1, image2), dim=0)
out = out[:number_of_images, :, :, :]
elif mode == "upto" and image2 is None:
out = image1.clone()
out = out[:number_of_images, :, :, :]
if number_of_repeats > 1:
out = out.repeat(number_of_repeats, 1, 1, 1)
# make sure to apply number of images
out = out[:number_of_images, :, :, :]
torch.cuda.empty_cache()
gc.collect()
return (out,)
@@ -1685,7 +1739,10 @@ class TextRender:
# For blank lines, add a blank line of the previous line's height
y_offset += previous_height + line_spacing
else:
w, h = d.textsize(line, font=font)
#w, h = d.textsize(line, font=font)
(left, top, right, bottom) = d.textbbox((0, 0), text=line, font=font)
w = right - left
h = bottom - top
previous_height = h # Update the previous height
if w > width:
requirements.txt
+3 -1
View File
@@ -4,4 +4,6 @@ google-auth-httplib2
oauth2client
httplib2
GitPython
librosa
librosa
moviepy
openai-whisper
video.py
+302 -88
View File
@@ -23,6 +23,8 @@ from skimage.metrics import structural_similarity as ssim
import cv2
import numpy as np
import cv2
import subprocess
import imageio_ffmpeg as ffmpeg
NODE_CLASS_MAPPINGS = {} # this is the dictionary that will be used to register the nodes
@@ -53,6 +55,17 @@ def video_files(s):
return ret
def run_ffmpeg_command(args):
"""
Executes a command using the ffmpeg executable.
Parameters:
args (list): List of command-line arguments for ffmpeg.
"""
cmd = [ffmpeg.get_ffmpeg_exe()] + args
subprocess.run(cmd)
# start with an ABC to define the common widget interface
class ABCWidgetMetaclass(ABCMeta):
@@ -276,22 +289,108 @@ class VideoFramesFolderToImageStack(metaclass=ABCWidgetMetaclass):
return (tensor,)
class SmoothStackTemporalByDistance(ABCABCVideoFolderToImage, metaclass=ABCWidgetMetaclass):
class SmoothStackTemporalByDistance2(ABCABCVideoFolderToImage, metaclass=ABCWidgetMetaclass):
@classmethod
def INPUT_TYPES(cls):
return {"required":
{
"image_stack": ("IMAGE",),
"top_threshold_first_order": ("FLOAT", {"min": 0.0, "max": 100.0, "step": 0.1, "default": 1.0}),
"bottom_threshold_first_order": ("FLOAT", {"min": 0.0, "max": 100.0, "step": 0.1, "default": 1.0}),
"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",)
RETURN_TYPES = ("IMAGE", "STRING")
def handler(self, image_stack,
top_threshold_first_order=1.0, bottom_threshold_first_order=1.0, iterations=1):
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(
@@ -301,38 +400,60 @@ class SmoothStackTemporalByDistance(ABCABCVideoFolderToImage, metaclass=ABCWidge
# 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)
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[:-2] > top_threshold_first_order)[0]
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 + 2])
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[:-2] < bottom_threshold_first_order)[0]
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]
# 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),)
return (image_stack.to(torch.float32), info,)
class GetImageStackStatisticsABC(ABCABCVideoFolderToImage, metaclass=ABCWidgetMetaclass):
class GetImageStackStatistics(ABCABCVideoFolderToImage, metaclass=ABCWidgetMetaclass):
"""Get statistics of an image stack"""
@classmethod
@@ -340,43 +461,54 @@ class GetImageStackStatisticsABC(ABCABCVideoFolderToImage, metaclass=ABCWidgetMe
return {"required":
{
"image_stack": ("IMAGE",),
"method": (["SSIM", "MSE"], {"default": "SSIM"}),
},
"optional":
{"text":
("STRING", {"multiline": False}),
}
}
RETURN_TYPES = ("STRING", "IMAGE",)
def handler(self, image_stack, text):
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 = image_stack.permute(0, 3, 1, 2)
image_stack_diff = image_stack[1:] - image_stack[:-1]
image_stack_diff = torch.sqrt(torch.sum(image_stack_diff ** 2, dim=(1, 2, 3)))
image_stack_diff_len = len(image_stack_diff)
required_width = 512 + image_stack_diff_len
image_stack_BCWH = image_stack.permute(0, 3, 1, 2)
image_stack_BCWH_DIFF = image_stack_BCWH[1:] - image_stack_BCWH[:-1]
mean = torch.mean(image_stack_diff)
std = torch.std(image_stack_diff)
median = torch.median(image_stack_diff)
max = torch.max(image_stack_diff)
min = torch.min(image_stack_diff)
# 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
out_str = f"mean: {mean}, std: {std}, median: {median}, max: {max}, min: {min}, len: {image_stack_diff_len}"
if method == "MSE":
stack_mse_B = torch.mean(image_stack_BCWH_DIFF ** 2, dim=(1, 2, 3))
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 = torch.max(stack_mse_B)
min = torch.min(stack_mse_B)
out_str = f"mean: {mean}\n std: {std}\n median: {median}\n max: {max}\n min: {min}\n len: {image_stack_diff_len}"
# create the histogram
hist_values, _ = np.histogram(image_stack_diff.numpy(), bins=100)
hist_values, _ = np.histogram(stack_mse_B.numpy(), bins=100)
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_width, 3))
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)
@@ -391,7 +523,7 @@ class GetImageStackStatisticsABC(ABCABCVideoFolderToImage, metaclass=ABCWidgetMe
# Create a plot for the difference for each frame pair
plot_offset = 512
for i, value in enumerate(image_stack_diff):
for i, value in enumerate(stack_mse_B):
height = int(value * scale_height / max) # Normalize by the max value
start_x = plot_offset + i
end_x = plot_offset + i + 1
@@ -457,6 +589,10 @@ class ImageStackToVideoFramesFolder(metaclass=ABCWidgetMetaclass):
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
@@ -467,9 +603,8 @@ class ImageStackToVideoFramesFolder(metaclass=ABCWidgetMetaclass):
# 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
import uuid
import shutil
random_folder = os.path.join(folder_out_path, str(uuid.uuid4()))
os.makedirs(random_folder)
for file in os.listdir(folder_out_path):
@@ -1014,12 +1149,18 @@ class ImageStackToVideoFile(metaclass=ABCWidgetMetaclass):
{
"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":
{"fps": ("FLOAT", {"default": 30, "min": 1, "max": 120}),
"func": ("FUNC", {"default": "NONE"}),
"audio_file": ("STRING", {"default": "NONE"}),
},
{
"audio_file": ("STRING", {"default": "NONE"}),
},
}
RETURN_TYPES = ("FUNC", "STRING",)
@@ -1029,75 +1170,106 @@ class ImageStackToVideoFile(metaclass=ABCWidgetMetaclass):
FUNCTION = "handler"
def handler(self, **kwargs):
"""frames are in the format (T,H,W,C)"""
import cv2
import os
import numpy as np
import uuid
import shutil
import inspect
import imageio
import audioread
import cv2
func = kwargs.get("func", None)
audio_file = kwargs.get("audio_file", None)
if func:
kwargs = func(**kwargs)
# video_dir = kwargs["video_dir"]
self.ARGS = kwargs
self.IN_FUNC = func
# get this functions source code
try:
self.CODE = inspect.getsource(ImageStackToVideoFile.handler)
except OSError:
self.CODE = "Unable to get source code"
audio_file = kwargs.get("audio_file")
if audio_file == "NONE":
audio_file = None
self.FUNC = self.handler
try:
image_stack = kwargs["image_stack"]
except KeyError:
# this means we are just passing our function through
return (self,)
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"]
fps = kwargs["fps"]
fps = np.array([fps], dtype=np.float16)[0]
video_out_path = os.path.join(video_dir, video_out)
# check to see if the video file already exists
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):
# move the existing file to a random folder inside the video folder
random_folder = os.path.join(video_dir, str(uuid.uuid4()))
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)
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
video = cv2.VideoWriter(video_out_path + ".mp4", fourcc, fps,
(image_stack.shape[2], image_stack.shape[1]))
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
for i in range(image_stack.shape[0]):
image = image_stack[i].numpy()
frames_length = image_stack.shape[0]
if fps_or_length == "fps":
video_length = frames_length / fps_length_value
else:
video_length = fps_length_value
# Check that the image is in the correct format
if "float16" in str(image.dtype):
image = image.astype(np.float32)
if fps_or_length == "fps":
video_fps = fps_length_value
else:
video_fps = frames_length / fps_length_value
if image.dtype != np.uint8:
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
image = (image * 255).astype(np.uint8) # convert to 0-255 and cast to uint8
folder_out = "temp_frames_folder"
video.write(image)
# Now, read the frames from disk as numpy data
import tempfile
# use tempfile to create a temporary folder
with tempfile.TemporaryDirectory() as temp_dir:
frames_out_path = temp_dir
if audio_file and os.path.exists(audio_file):
video_with_audio_path = video_out_path + "_with_audio.mp4"
cmd = f"ffmpeg -i {video_out_path}.mp4 -i {audio_file} -filter_complex \
'[0:a]apad=pad_dur={video_duration}[a];[a][1:a]amerge=inputs=2[aout]' \
-map [aout] -c:v copy -c:a aac {video_with_audio_path}"
os.system(cmd)
return (self, video_with_audio_path)
video.release()
image_stack = None
del image_stack
return (self, video_out_path + ".mp4",)
# Convert the image stack to frames on disk
frame_folder_handler = ImageStackToVideoFramesFolder()
frame_folder_handler.handler(image_stack=image_stack, folder_out=frames_out_path)
frames = []
for fl in sorted(os.listdir(frames_out_path)):
# if the fl is an actual image file
if not fl.lower().endswith(('.jpg', '.jpeg', '.png')):
continue
# create the numpy array
frame = cv2.imread(os.path.join(frames_out_path, fl))
# convert to RGB
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
# append to the list of frames
frames.append(frame)
writer = imageio.get_writer(video_out_path, fps=video_fps)
for frame in frames:
writer.append_data(frame)
writer.close()
if audio_file is not None:
if audio_length > video_length:
while video_length < audio_length:
image_stack = np.concatenate((image_stack, image_stack[-1:]), axis=0)
video_length += 1 / video_fps
frames_length += 1
args = [
"-i", video_out_path,
"-i", audio_file,
"-c:v", codec,
"-b:v", str(bitrate),
"-c:a", "aac",
"-strict", "experimental",
video_out_path
]
run_ffmpeg_command(args)
return (self, video_out_path,)
class CreateConsistentVideo(metaclass=ABCWidgetMetaclass):
@@ -1285,3 +1457,45 @@ def combine_videos(video_paths, base_input_directory, output_path):
# Clean up the temporary file
temp_file.close()
os.remove(temp_file.name)
class WhisperTranscription(metaclass=ABCWidgetMetaclass):
@classmethod
def INPUT_TYPES(cls):
return {"required": {"file_path": ("STRING", {"default": "fullpath"})}}
RETURN_TYPES = ("STRING",)
FUNCTION = "transcribe_movie"
CATEGORY = "transcription"
def transcribe_movie(self, file_path):
import os
import time
import numpy as np
import whisper as wh
import moviepy as mp
try:
model = wh.load_model("medium.en")
filename = os.path.basename(file_path)
audio_file_path = f"audio_{filename}.wav"
video = mp.VideoFileClip(file_path)
video.audio.write_audiofile(audio_file_path)
res = wh.transcribe(model=model, audio=audio_file_path)
segs = res["segments"]
out = []
for seg in segs:
out.append(f"{seg['start']:3.3f} - {seg['end']:3.3f}: {seg['text']}")
ret = '\n'.join(out)
os.remove(audio_file_path)
return (ret,)
except Exception as e:
print(e)
return ("Error",)