From 5669d26274eb53c9fa263fe28b16faa327928d85 Mon Sep 17 00:00:00 2001
From: filipstrand <strand.filip@gmail.com>
Date: Sun, 18 Aug 2024 20:41:48 +0200
Subject: [PATCH] Never override previously generated images

---
 README.md                                     |  3 ++-
 main.py                                       |  3 ++-
 src/flux_1_schnell/config/config.py           |  1 +
 .../post_processing/image_util.py             | 25 +++++++++++++++++++
 4 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 1d1d1f8..06bebbe 100644
--- a/README.md
+++ b/README.md
@@ -51,6 +51,7 @@ sys.path.append("/path/to/mflux/src")
 
 from flux_1_schnell.config.config import Config
 from flux_1_schnell.flux import Flux1Schnell
+from flux_1_schnell.post_processing.image_util import ImageUtil
 
 flux = Flux1Schnell("black-forest-labs/FLUX.1-schnell")
 
@@ -62,7 +63,7 @@ image = flux.generate_image(
    )
 )
 
-image.save("image.png")
+ImageUtil.save_image(image, "image.png")
 ```
 
 If the model is not already downloaded on your machine, it will start the download process and fetch the model weights (~34GB in size for the Schnell model).
diff --git a/main.py b/main.py
index ba639a1..5cf39bb 100644
--- a/main.py
+++ b/main.py
@@ -5,6 +5,7 @@ sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'src')))
 
 from flux_1_schnell.config.config import Config
 from flux_1_schnell.flux import Flux1Schnell
+from flux_1_schnell.post_processing.image_util import ImageUtil
 
 flux = Flux1Schnell("black-forest-labs/FLUX.1-schnell")
 
@@ -18,4 +19,4 @@ image = flux.generate_image(
     )
 )
 
-image.save("image.png")
+ImageUtil.save_image(image, "image.png")
diff --git a/src/flux_1_schnell/config/config.py b/src/flux_1_schnell/config/config.py
index 2a01fd2..51e9599 100644
--- a/src/flux_1_schnell/config/config.py
+++ b/src/flux_1_schnell/config/config.py
@@ -4,6 +4,7 @@ import logging
 
 log = logging.getLogger(__name__)
 
+
 class Config:
     precision: mx.Dtype = mx.float16
     num_train_steps = 1000
diff --git a/src/flux_1_schnell/post_processing/image_util.py b/src/flux_1_schnell/post_processing/image_util.py
index 95f175a..9cc282a 100644
--- a/src/flux_1_schnell/post_processing/image_util.py
+++ b/src/flux_1_schnell/post_processing/image_util.py
@@ -1,8 +1,13 @@
+import logging
+from pathlib import Path
+
 import PIL
 import mlx.core as mx
 import numpy as np
 from PIL import Image
 
+log = logging.getLogger(__name__)
+
 
 class ImageUtil:
 
@@ -53,3 +58,23 @@ class ImageUtil:
     def resize(image):
         image = image.resize((1024, 1024), resample=PIL.Image.LANCZOS)
         return image
+
+    @staticmethod
+    def save_image(image: Image.Image, path: str) -> None:
+        file_path = Path(path)
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+        file_name = file_path.stem
+        file_extension = file_path.suffix
+
+        # If a file already exists, create a new name with a counter
+        counter = 1
+        while file_path.exists():
+            new_name = f"{file_name}({counter}){file_extension}"
+            file_path = file_path.with_name(new_name)
+            counter += 1
+
+        try:
+            image.save(file_path)
+            log.info(f"Image saved successfully at: {file_path}")
+        except Exception as e:
+            log.info(f"Error saving image: {e}")