Save in parallel

nodes.py

@@ -707,6 +707,8 @@ class SaveImage:
     CATEGORY = "image"
 
     def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
+        from joblib import Parallel, delayed
+
         def map_filename(filename):
             prefix_len = len(filename_prefix)
             prefix = filename[:prefix_len + 1]
@@ -715,16 +717,8 @@ class SaveImage:
             except:
                 digits = 0
             return (digits, prefix)
-        try:
-            counter = max(filter(lambda a: a[1][:-1] == filename_prefix and a[1][-1] == "_", map(map_filename, os.listdir(self.output_dir))))[0] + 1
-        except ValueError:
-            counter = 1
-        except FileNotFoundError:
-            os.mkdir(self.output_dir)
-            counter = 1
-
-        paths = list()
-        for image in images:
+    
+        def save_image(image, file):
             i = 255. * image.cpu().numpy()
             img = Image.fromarray(i.astype(np.uint8))
             metadata = PngInfo()
@@ -733,10 +727,19 @@ class SaveImage:
             if extra_pnginfo is not None:
                 for x in extra_pnginfo:
                     metadata.add_text(x, json.dumps(extra_pnginfo[x]))
-            file = f"{filename_prefix}_{counter:05}_.png"
             img.save(os.path.join(self.output_dir, file), pnginfo=metadata, optimize=True)
-            paths.append(file)
-            counter += 1
+
+        try:
+            counter = max(filter(lambda a: a[1][:-1] == filename_prefix and a[1][-1] == "_", map(map_filename, os.listdir(self.output_dir))))[0] + 1
+        except ValueError:
+            counter = 1
+        except FileNotFoundError:
+            os.mkdir(self.output_dir)
+            counter = 1
+
+        paths = [f"{filename_prefix}_{i:05}_.png" for i in range(counter, counter + len(images))]
+        parallelism = min(4, len(images))
+        Parallel(n_jobs=parallelism)(delayed(save_image)(img, path) for (img, path) in zip(images, paths))
         return { "ui": { "images": paths } }
 
 class LoadImage:

requirements.txt

@@ -9,3 +9,4 @@ safetensors
 pytorch_lightning
 aiohttp
 accelerate
+joblib