Merge pull request #132 from filipstrand/add-option-to-override-saved-images

Add option to override saved images

src/mflux/post_processing/generated_image.py

@@ -44,10 +44,15 @@ class GeneratedImage:
         self.init_image_path = init_image_path
         self.init_image_strength = init_image_strength
 
-    def save(self, path: t.Union[str, pathlib.Path], export_json_metadata: bool = False) -> None:
+    def save(
+        self,
+        path: t.Union[str, pathlib.Path],
+        export_json_metadata: bool = False,
+        overwrite: bool = False,
+    ) -> None:
         from mflux import ImageUtil
 
-        ImageUtil.save_image(self.image, path, self._get_metadata(), export_json_metadata)
+        ImageUtil.save_image(self.image, path, self._get_metadata(), export_json_metadata, overwrite)
 
     def _get_metadata(self) -> dict:
         """Generate metadata for reference as well as input data for

src/mflux/post_processing/image_util.py

@@ -114,22 +114,24 @@ class ImageUtil:
 
     @staticmethod
     def save_image(
-            image: PIL.Image.Image,
-            path: t.Union[str, pathlib.Path],
-            metadata: dict | None = None,
-            export_json_metadata: bool = False
+        image: PIL.Image.Image,
+        path: t.Union[str, pathlib.Path],
+        metadata: dict | None = None,
+        export_json_metadata: bool = False,
+        overwrite: bool = False
     ) -> None:  # fmt: off
         file_path = pathlib.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
+        # If a file already exists and overwrite is False, create a new name with a counter
+        if not overwrite:
+            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:
             # Save image without metadata first

tests/image_generation/helpers/image_generation_controlnet_test_helper.py

@@ -48,7 +48,7 @@ class ImageGeneratorControlnetTestHelper:
                     controlnet_strength=controlnet_strength,
                 ),
             )
-            image.save(path=output_image_path)
+            image.save(path=output_image_path, overwrite=True)
 
             # then
             np.testing.assert_array_equal(

tests/image_generation/helpers/image_generation_test_helper.py

@@ -48,7 +48,7 @@ class ImageGeneratorTestHelper:
                     width=width,
                 ),
             )
-            image.save(path=output_image_path)
+            image.save(path=output_image_path, overwrite=True)
 
             # then
             np.testing.assert_array_equal(