Merge pull request #9 from filipstrand/never-override-previously-generated-images

Never override previously generated images

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).

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")

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

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}")