From aed1618bf74473ccd7c2d03ecdda0d63efe50062 Mon Sep 17 00:00:00 2001 From: Tasha Upchurch Date: Thu, 7 Sep 2023 22:02:39 -0400 Subject: [PATCH] added whisper-openai to requirements as well as moviepy 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. --- basic.py | 179 ++++++++++++++++++++++ image.py | 157 +++++++++++++------ requirements.txt | 4 +- video.py | 390 ++++++++++++++++++++++++++++++++++++----------- 4 files changed, 591 insertions(+), 139 deletions(-) create mode 100644 basic.py diff --git a/basic.py b/basic.py new file mode 100644 index 0000000..39d2db3 --- /dev/null +++ b/basic.py @@ -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 diff --git a/image.py b/image.py index f6a2645..a4b04f2 100644 --- a/image.py +++ b/image.py @@ -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: diff --git a/requirements.txt b/requirements.txt index 7364b1d..dbcc08e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,4 +4,6 @@ google-auth-httplib2 oauth2client httplib2 GitPython -librosa \ No newline at end of file +librosa +moviepy +openai-whisper \ No newline at end of file diff --git a/video.py b/video.py index 0463eea..9d85a30 100644 --- a/video.py +++ b/video.py @@ -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",) \ No newline at end of file