added initial scaffolding for SaveImageSequence node

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,27 +376,87 @@ 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:
             m.update(f.read())
         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 🎥🅥🅗🅢",
 }

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;
+			};
+		}
+    },
+});

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);

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]
 		}
 	},
 });