Add files via upload

nodes/FV_SaveSimpleJPG.py

@@ -2,9 +2,8 @@ import numpy as np
 from PIL import Image
 import folder_paths
 import os
-import re
 
-class SaveSimpleJPG:
+class Save_as_jpg:
     def __init__(self):
         self.output_dir = folder_paths.get_output_directory()
         self.type = "output"
@@ -15,51 +14,70 @@ class SaveSimpleJPG:
             "required": {
                 "images": ("IMAGE",),
                 "filename_prefix": ("STRING", {"default": "ComfyUI"}),
-                "quality": ("INT", {"default": 80, "min": 10, "max": 100, "step": 1}),
+                "compression": ("INT", {"default": 80, "min": 10, "max": 100, "step": 1}),
             }
         }
 
     RETURN_TYPES = ()
-    FUNCTION = "save_jpg"
+    FUNCTION = "Save_as_jpg"
     OUTPUT_NODE = True
     CATEGORY = "Fictiverse"
 
-    def get_next_counter(self, folder, prefix):
-        if not os.path.exists(folder):
-            return 1
+    def Save_as_jpg(self, images, filename_prefix="ComfyUI", compression=80):
+        def map_filename(filename):
+            prefix_len = len(os.path.basename(filename_prefix))
+            prefix = filename[:prefix_len + 1]
+            try:
+                digits = int(filename[prefix_len + 1:].split('_')[0])
+            except:
+                digits = 0
+            return (digits, prefix)
 
-        pattern = re.compile(rf"^{re.escape(prefix)}_(\d+)\.jpg$")
-        max_id = 0
+        def compute_vars(input):
+            input = input.replace("%width%", str(images[0].shape[1]))
+            input = input.replace("%height%", str(images[0].shape[0]))
+            return input
 
-        for f in os.listdir(folder):
-            match = pattern.match(f)
-            if match:
-                num = int(match.group(1))
-                if num > max_id:
-                    max_id = num
+        filename_prefix = compute_vars(filename_prefix)
 
-        return max_id + 1
+        subfolder = os.path.dirname(os.path.normpath(filename_prefix))
+        filename = os.path.basename(os.path.normpath(filename_prefix))
 
-    def save_jpg(self, images, filename_prefix="ComfyUI", quality=80):
-        output_folder = self.output_dir
-        os.makedirs(output_folder, exist_ok=True)
+        full_output_folder = os.path.join(self.output_dir, subfolder)
 
-        counter = self.get_next_counter(output_folder, filename_prefix)
+        if os.path.commonpath((self.output_dir, os.path.abspath(full_output_folder))) != self.output_dir:
+            print("Saving image outside the output folder is not allowed.")
+            return {}
 
-        results = []
+        try:
+            counter = max(
+                filter(
+                    lambda a: a[1][:-1] == filename and a[1][-1] == "_",
+                    map(map_filename, os.listdir(full_output_folder))
+                )
+            )[0] + 1
+        except ValueError:
+            counter = 1
+        except FileNotFoundError:
+            os.makedirs(full_output_folder, exist_ok=True)
+            counter = 1
+
+        results = list()
 
         for image in images:
-            i = 255. * image.cpu().numpy()
+            i = 255.0 * image.cpu().numpy()
             img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
 
-            filename = f"{filename_prefix}_{counter:05}.jpg"
-            filepath = os.path.join(output_folder, filename)
+            file = f"{filename}_{counter:05}_.jpg"
 
-            img.save(filepath, format="JPEG", quality=quality)
+            img.save(
+                os.path.join(full_output_folder, file),
+                quality=compression
+            )
 
             results.append({
-                "filename": filename,
-                "subfolder": "",
+                "filename": file,
+                "subfolder": subfolder,
                 "type": self.type
             })
 
@@ -69,5 +87,5 @@ class SaveSimpleJPG:
 
 
 NODE_CLASS_MAPPINGS = {
-    "Save Simple JPG": SaveSimpleJPG,
+    "Save_as_jpg": Save_as_jpg
 }