diff --git a/nodes.py b/nodes.py index 26dad5729..ab4201037 100644 --- a/nodes.py +++ b/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: diff --git a/requirements.txt b/requirements.txt index 45f2599d9..5fd675948 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,3 +9,4 @@ safetensors pytorch_lightning aiohttp accelerate +joblib