From 0a5bedf4124ac396bccae721e1f29f54e8b0e00a Mon Sep 17 00:00:00 2001
From: Rhys Yang <receyuki@gmail.com>
Date: Wed, 1 Nov 2023 04:22:07 +0800
Subject: [PATCH] Fix `%counter` overwriting existing images #11 #14, and add a
 new file-naming mechanism to ensure naming uniqueness.

---
 nodes.py | 37 ++++++++++++++++++++++++++++++-------
 1 file changed, 30 insertions(+), 7 deletions(-)

diff --git a/nodes.py b/nodes.py
index 226efa6..e7c28c7 100644
--- a/nodes.py
+++ b/nodes.py
@@ -8,6 +8,7 @@
 
 import os
 from datetime import datetime
+from itertools import chain
 
 import torch
 import json
@@ -293,7 +294,7 @@ class SDPromptSaver:
         (
             full_output_folder,
             filename_alt,
-            counter,
+            counter_alt,
             subfolder_alt,
             filename_prefix,
         ) = folder_paths.get_save_image_path(
@@ -316,7 +317,6 @@ class SDPromptSaver:
             variable_map = {
                 "%date": self.get_time(date_format),
                 "%time": self.get_time(time_format),
-                "%counter": f"{counter:05}",
                 "%seed": seed,
                 "%steps": steps,
                 "%cfg": cfg,
@@ -327,6 +327,12 @@ class SDPromptSaver:
                 "%quality": jpg_webp_quality,
             }
 
+            subfolder = self.get_path(path, variable_map)
+            output_folder = Path(full_output_folder) / subfolder
+            output_folder.mkdir(parents=True, exist_ok=True)
+            counter = self.get_counter(output_folder)
+            variable_map["%counter"] = f"{counter:05}"
+
             i = 255.0 * image.cpu().numpy()
             img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
             metadata = None
@@ -348,10 +354,9 @@ class SDPromptSaver:
                 f"Version: ComfyUI"
                 f"{extra_info_real}"
             )
-            subfolder = self.get_path(path, variable_map)
-            output_folder = Path(full_output_folder) / subfolder
-            output_folder.mkdir(parents=True, exist_ok=True)
-            file = self.get_path(filename, variable_map).with_suffix("." + extension)
+
+            stem = self.get_path(filename, variable_map)
+            file = self.get_unique_filename(stem, extension, output_folder)
 
             if extension == "png":
                 if not args.disable_metadata:
@@ -387,7 +392,6 @@ class SDPromptSaver:
             results.append(
                 {"filename": file.name, "subfolder": str(subfolder), "type": self.type}
             )
-            counter += 1
 
         return {"ui": {"images": results}}
 
@@ -403,6 +407,13 @@ class SDPromptSaver:
 
         return hash_sha256.hexdigest()[:10]
 
+    @staticmethod
+    def get_counter(directory: Path):
+        img_files = list(
+            chain(*(directory.rglob(f"*{suffix}") for suffix in SUPPORTED_FORMATS))
+        )
+        return len(img_files) + 1
+
     @staticmethod
     def get_path(name, variable_map):
         for variable, value in variable_map.items():
@@ -418,6 +429,18 @@ class SDPromptSaver:
         except:
             return ""
 
+    @staticmethod
+    def get_unique_filename(stem: Path, extension: str, output_folder: Path):
+        file = stem.with_suffix(f".{extension}")
+        index = 0
+
+        while (output_folder / file).exists():
+            index += 1
+            new_stem = f"{stem}_{index}"
+            file = Path(new_stem).with_suffix(f".{extension}")
+
+        return file
+
 
 class SDParameterGenerator:
     ASPECT_RATIO_MAP = {