diff --git a/videohelpersuite/nodes.py b/videohelpersuite/nodes.py index 673fb8a..cac15d7 100755 --- a/videohelpersuite/nodes.py +++ b/videohelpersuite/nodes.py @@ -212,15 +212,22 @@ class LoadVideo: @classmethod def INPUT_TYPES(s): video_extensions = ['webm', 'mp4', 'mkv', 'gif'] - input_dir = folder_paths.get_input_directory() - files = [] - for f in os.listdir(input_dir): - if os.path.isfile(os.path.join(input_dir, f)): - file_parts = f.split('.') - if len(file_parts) > 1 and (file_parts[-1] in video_extensions): - files.append(f) + input_dirs = [folder_paths.get_input_directory(), folder_paths.get_output_directory(), folder_paths.get_temp_directory()] + files_list = [] + for source in input_dirs: + files = [] + if not os.path.exists(source): + files_list.append(files) + continue + for f in os.listdir(source): + if os.path.isfile(os.path.join(source, f)): + file_parts = f.split('.') + if len(file_parts) > 1 and (file_parts[-1] in video_extensions): + files.append(f) + files_list.append(sorted(files)) return {"required": { - "video": (sorted(files), {"video_upload": True}), + "video": (files_list, {"video_upload": True}), + "search_folder": (["input","output","temp"],), "force_rate": ("INT", {"default": 0, "min": 0, "max": 24, "step": 1}), "force_size": (["Disabled", "256x?", "?x256", "256x256", "512x?", "?x512", "512x512"],), "frame_load_cap": ("INT", {"default": 0, "min": 0, "step": 1}), @@ -254,9 +261,9 @@ def target_size(self, width, height, force_size): width = int(force_size[0]) return (width, height) - def load_video_cv_fallback(self, video, force_rate, force_size, frame_load_cap, skip_first_frames): + def load_video_cv_fallback(self, video, search_folder, force_rate, force_size, frame_load_cap, skip_first_frames): try: - video_cap = cv2.VideoCapture(folder_paths.get_annotated_filepath(video)) + video_cap = cv2.VideoCapture(folder_paths.get_annotated_filepath(video, default_dir=search_folder)) if not video_cap.isOpened(): raise ValueError(f"{video} could not be loaded with cv fallback.") # set video_cap to look at start_index frame @@ -312,13 +319,13 @@ def load_video_cv_fallback(self, video, force_rate, force_size, frame_load_cap, # TODO: raise an error maybe if no frames were loaded? return (images, frames_added) - def load_video(self, video, force_rate, force_size, frame_load_cap, skip_first_frames): + def load_video(self, video, search_folder, force_rate, force_size, frame_load_cap, skip_first_frames): # check if video is a gif - will need to use cv fallback to read frames # use cv fallback if ffmpeg not installed or gif if ffmpeg_path is None: - return self.load_video_cv_fallback(video, force_rate, force_size, frame_load_cap, skip_first_frames) + return self.load_video_cv_fallback(video, search_folder, force_rate, force_size, frame_load_cap, skip_first_frames) # otherwise, continue with ffmpeg - video_path = folder_paths.get_annotated_filepath(video) + video_path = folder_paths.get_annotated_filepath(video, default_dir=search_folder) args_dummy = [ffmpeg_path, "-i", video_path, "-f", "null", "-"] try: with subprocess.Popen(args_dummy, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE) as proc: @@ -329,7 +336,7 @@ def load_video(self, video, force_rate, force_size, frame_load_cap, skip_first_f break except Exception as e: logger.info(f"Retrying with opencv due to ffmpeg error: {e}") - return self.load_video_cv_fallback(video, force_rate, force_size, frame_load_cap, skip_first_frames) + return self.load_video_cv_fallback(video, search_folder, force_rate, force_size, frame_load_cap, skip_first_frames) args_all_frames = [ffmpeg_path, "-i", video_path, "-v", "error", "-pix_fmt", "rgb24"] @@ -369,13 +376,13 @@ def load_video(self, video, force_rate, force_size, frame_load_cap, skip_first_f current_offset = 0 except Exception as e: logger.info(f"Retrying with opencv due to ffmpeg error: {e}") - return self.load_video_cv_fallback(video, force_rate, force_size, frame_load_cap, skip_first_frames) + return self.load_video_cv_fallback(video, search_folder, force_rate, force_size, frame_load_cap, skip_first_frames) images = torch.from_numpy(np.stack(images)) return (images, images.size(0)) @classmethod - def IS_CHANGED(s, video, force_size, frame_load_cap, skip_first_frames): + def IS_CHANGED(s, video, **kwargs): image_path = folder_paths.get_annotated_filepath(video) m = hashlib.sha256() with open(image_path, 'rb') as f: @@ -383,13 +390,73 @@ def IS_CHANGED(s, video, force_size, frame_load_cap, skip_first_frames): return m.digest().hex() @classmethod - def VALIDATE_INPUTS(s, video, **kwargs): - if not folder_paths.exists_annotated_filepath(video): + def VALIDATE_INPUTS(s, video, search_folder, **kwargs): + if not folder_paths.exists_annotated_filepath(video + f" [{search_folder}]"): return "Invalid image file: {}".format(video) return True +class SaveImageSequence: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE",), + "timestamp_directory": ("BOOLEAN", {"default": False}), + "directory_name": ("STRING", {"default": "AnimateDiff"}), + "filename_prefix": ("STRING", {"default": "image"}), + "starting_number": ("INT", {"default": 0, "min": 0, "step": 1}), + "padding": ("INT", {"default": 0, "min": 0, "step": 1}), + }, + "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, + } + + CATEGORY = "Video Helper Suite 🎥🅥🅗🅢" + + OUTPUT_NODE = True + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("directory",) + FUNCTION = "save_images" + + def save_images(self, images: torch.Tensor, timestamp_directory, directory_name, + filename_prefix, starting_number, padding, prompt = None, extra_pnginfo=None): + #NOTE: timestamp_directory is an unused placeholder for javascript + # If set, client code will automatically set the directory name + # goal: save image sequence, + # allowing user to choose padding amount and starting fraame number + # as well as subdirectory in output to save it as. + # Output directory should have option to either be a set name, or be dynamically + # generated with timestamp. + output_dir = os.path.join(folder_paths.get_output_directory(), directory_name) + os.makedirs(output_dir, exist_ok=True) + metadata = PngInfo() + if prompt is not None: + metadata.add_text("prompt", json.dumps(prompt)) + if extra_pnginfo is not None: + for x in extra_pnginfo: + metadata.add_text(x, json.dumps(extra_pnginfo[x])) + + for i, image in enumerate(images): + img = 255.0 * image.cpu().numpy() + img = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8)) + file = f"{filename_prefix}{{0:0{padding}d}}.png".format(i+starting_number) + file_path = os.path.join(output_dir, file) + img.save(file_path, pnginfo=metadata, compress_level=4) + + return output_dir + + + @classmethod + def VALIDATE_INPUTS(s, **kwargs): + if kwargs['directory_name'].startswith('/') or '..' in kwargs['directory_name']: + return "Invalid, potentially dangerous directory name: " + kwargs['directory_name'] + if kwargs['filename_prefix'].startswith('/') or '..' in kwargs['filename_prefix']: + return "Invalid, potentially dangerous prefix name: " + kwargs['filename_prefix'] + + return True + + class LoadImagesFromDirectory: @classmethod def INPUT_TYPES(s): @@ -495,6 +562,7 @@ def VALIDATE_INPUTS(s, directory: str, **kwargs): "VHS_GetImageCount": GetImageCount, "VHS_DuplicateLatents": DuplicateLatents, "VHS_DuplicateImages": DuplicateImages, + "VHS_SaveImageSequence": SaveImageSequence, } NODE_DISPLAY_NAME_MAPPINGS = { "VHS_VideoCombine": "Video Combine 🎥🅥🅗🅢", @@ -511,4 +579,5 @@ def VALIDATE_INPUTS(s, directory: str, **kwargs): "VHS_GetImageCount": "Get Image Count 🎥🅥🅗🅢", "VHS_DuplicateLatents": "Duplicate Latent Batch 🎥🅥🅗🅢", "VHS_DuplicateImages": "Duplicate Image Batch 🎥🅥🅗🅢", + "VHS_SaveImageSequence": "Save Image Sequence 🎥🅥🅗🅢", } diff --git a/web/js/dateformatting.js b/web/js/dateformatting.js new file mode 100644 index 0000000..026c163 --- /dev/null +++ b/web/js/dateformatting.js @@ -0,0 +1,81 @@ +import { app } from '../../../scripts/app.js' + +// Simple date formatter +const parts = { + d: (d) => d.getDate(), + M: (d) => d.getMonth() + 1, + h: (d) => d.getHours(), + m: (d) => d.getMinutes(), + s: (d) => d.getSeconds(), +}; +const format = + Object.keys(parts) + .map((k) => k + k + "?") + .join("|") + "|yyy?y?"; + +function formatDate(text, date) { + return text.replace(new RegExp(format, "g"), function (text) { + if (text === "yy") return (date.getFullYear() + "").substring(2); + if (text === "yyyy") return date.getFullYear(); + if (text[0] in parts) { + const p = parts[text[0]](date); + return (p + "").padStart(text.length, "0"); + } + return text; + }); +} + +app.registerExtension({ + name: "VideoHelperSuite.DateFormatting", + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if (nodeData?.name == "VHS_VideoCombine") { + const onNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + const r = onNodeCreated ? onNodeCreated.apply(this, arguments) : undefined; + + const widget = this.widgets.find((w) => w.name === "filename_prefix"); + widget.serializeValue = () => { + return widget.value.replace(/%([^%]+)%/g, function (match, text) { + const split = text.split("."); + if (split[0].startsWith("date:")) { + return formatDate(split[0].substring(5), new Date()); + } + return match; + }); + }; + + return r; + }; + } + if (nodeData?.name == "VHS_SaveImageSequence") { + const onNodeCreated = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + const r = onNodeCreated ? onNodeCreated.apply(this, arguments) : undefined; + + const directoryWidget = this.widgets.find((w) => w.name === "directory_name"); + const timestampWidget = this.widgets.find((w) => w.name === "timestamp_directory"); + directoryWidget.serializeValue = () => { + if (timestampWidget.value) { + //ignore actual value and return timestamp + return formatDate("yyyy-MM-ddThh:mm:ss", new Date()); + } + return directoryWidget.value + }; + directoryWidget._value = directoryWidget.value; + Object.defineProperty(directoryWidget, "value", { + set : function(value) { + directoryWidget._value = value; + }, + get : function() { + if (timestampWidget.value) { + return "yyyy-MM-ddThh:mm:ss"; + } + return directoryWidget._value; + } + }); + + return r; + }; + } + }, +}); diff --git a/web/js/folderupload.js b/web/js/folderupload.js index dbc6da5..28014cd 100644 --- a/web/js/folderupload.js +++ b/web/js/folderupload.js @@ -54,7 +54,6 @@ function folderUpload(node, inputName, inputData, app) { if (i > 0) { directoryWidget.value = directory.slice(0,directory.lastIndexOf('/')) } - app.file = fileInput.files[0] }, }); document.body.append(fileInput); diff --git a/web/js/videoupload.js b/web/js/videoupload.js index f4d4fb2..1faea3f 100644 --- a/web/js/videoupload.js +++ b/web/js/videoupload.js @@ -3,12 +3,47 @@ import { api } from '../../../scripts/api.js' import { ComfyWidgets } from "../../../scripts/widgets.js" function videoUpload(node, inputName, inputData, app) { - const imageWidget = node.widgets.find((w) => w.name === "video"); + const pathWidget = node.widgets.find((w) => w.name === "video"); + const folderWidget = node.widgets.find((w) => w.name === "search_folder"); + + folderWidget.content_index = 0; let uploadWidget; + pathWidget.contentlists = inputData[1] + Object.defineProperty(folderWidget, "value", { + set : function(value) { + this.content_index = this.options.values.indexOf(value); + if (this.content_index == -1) { + this.content_index = 0 + } + pathWidget.options.values = pathWidget.contentlists[this.content_index]; + if (pathWidget.options.values.length == 0) { + pathWidget.value = "None"; + } else { + pathWidget.value = pathWidget.options.values[0]; + } + this._value = value; + }, + get : function() { + return this._value; + } + }); - var default_value = imageWidget.value; - Object.defineProperty(imageWidget, "value", { + var default_value = "None"; + if (inputData[1][0].length > 0) { + default_value = inputData[1][0][0]; + } + Object.defineProperty(pathWidget, "value", { set : function(value) { + if (typeof(value) == 'object') { + //refresh event + this.contentlists = this.options.values; + this.options.values = this.contentlists[folderWidget.content_index]; + if (this.options.values.length > 0) { + value = this.options.values[0]; + } else { + value = "None"; + } + } this._real_value = value; }, @@ -52,12 +87,12 @@ function videoUpload(node, inputName, inputData, app) { let path = data.name; if (data.subfolder) path = data.subfolder + "/" + path; - if (!imageWidget.options.values.includes(path)) { - imageWidget.options.values.push(path); + if (!pathWidget.options.values.includes(path)) { + pathWidget.options.values.push(path); } if (updateNode) { - imageWidget.value = path; + pathWidget.value = path; } } else { alert(resp.status + " - " + resp.statusText); @@ -90,11 +125,11 @@ function videoUpload(node, inputName, inputData, app) { ComfyWidgets.VIDEOUPLOAD = videoUpload; app.registerExtension({ - name: "VideoHelperSuit.UploadVideo", + name: "VideoHelperSuite.UploadVideo", async beforeRegisterNodeDef(nodeType, nodeData, app) { if (nodeData?.name == "VHS_LoadVideo") { - console.log("test") - nodeData.input.required.upload = ["VIDEOUPLOAD"]; + nodeData.input.required.upload = ["VIDEOUPLOAD", nodeData.input.required.video[0]]; + nodeData.input.required.video[1] = nodeData.input.required.video[1][0] } }, });