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