Merge branch 'filipstrand:main' into memory_saver

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/dreambooth/test_resume_training.py

@@ -14,6 +14,9 @@ CHECKPOINT_5 = "tests/dreambooth/tmp/_checkpoints/0000005_checkpoint.zip"
 
 class TestResumeTraining:
     def test_resume_training(self):
+        # Clean up any existing temporary directories from previous test runs
+        TestResumeTraining.delete_folder_if_exists("tests/dreambooth/tmp")
+
         try:
             # Given: A small training run from scratch for 5 steps (as described in the config)...
             fluxA, runtime_config, training_spec, training_state = DreamBoothInitializer.initialize(
@@ -85,3 +88,11 @@ class TestResumeTraining:
     @staticmethod
     def delete_folder(path: str) -> None:
         return shutil.rmtree(path)
+
+    @staticmethod
+    def delete_folder_if_exists(path: str) -> None:
+        if os.path.exists(path):
+            shutil.rmtree(path)
+            print(f"Deleted folder: {path}")
+        else:
+            print("The specified folder does not exist.")

tests/dreambooth/test_train_and_load_weights.py

@@ -1,3 +1,4 @@
+import os
 import shutil
 
 import numpy as np
@@ -14,6 +15,9 @@ LORA_FILE = "tests/dreambooth/tmp/_checkpoints/0000005_checkpoint/0000005_adapte
 
 class TestTrainAndLoadWeights:
     def test_train_and_load_weights(self):
+        # Clean up any existing temporary directories from previous test runs
+        TestTrainAndLoadWeights.delete_folder_if_exists("tests/dreambooth/tmp")
+
         try:
             # Given: A small training run from scratch for 5 steps (as described in the config)...
             fluxA, runtime_config, training_spec, training_state = DreamBoothInitializer.initialize(
@@ -72,3 +76,11 @@ class TestTrainAndLoadWeights:
     @staticmethod
     def delete_folder(path: str) -> None:
         return shutil.rmtree(path)
+
+    @staticmethod
+    def delete_folder_if_exists(path: str) -> None:
+        if os.path.exists(path):
+            shutil.rmtree(path)
+            print(f"Deleted folder: {path}")
+        else:
+            print("The specified folder does not exist.")

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(

tests/model_saving/test_model_saving.py

@@ -1,3 +1,4 @@
+import os
 import shutil
 
 import numpy as np
@@ -9,6 +10,9 @@ PATH = "tests/4bit/"
 
 class TestModelSaving:
     def test_save_and_load_4bit_model(self):
+        # Clean up any existing temporary directories from previous test runs
+        TestModelSaving.delete_folder_if_exists(PATH)
+
         try:
             # given a saved quantized model (and an image from that model)
             fluxA = Flux1(
@@ -56,3 +60,11 @@ class TestModelSaving:
     @staticmethod
     def delete_folder(path: str) -> None:
         return shutil.rmtree(path)
+
+    @staticmethod
+    def delete_folder_if_exists(path: str) -> None:
+        if os.path.exists(path):
+            shutil.rmtree(path)
+            print(f"Deleted folder: {path}")
+        else:
+            print(f"Folder does not exist: {path}")

tests/model_saving/test_model_saving_lora.py

@@ -1,3 +1,4 @@
+import os
 import shutil
 from pathlib import Path
 
@@ -10,6 +11,9 @@ PATH = "tests/4bit/"
 
 class TestModelSavingLora:
     def test_save_and_load_4bit_model_with_lora(self):
+        # Clean up any existing temporary directories from previous test runs
+        TestModelSavingLora.delete_folder_if_exists(PATH)
+
         try:
             # given a saved quantized model on disk (without LoRA)...
             fluxA = Flux1(
@@ -71,6 +75,14 @@ class TestModelSavingLora:
     def delete_folder(path: str) -> None:
         return shutil.rmtree(path)
 
+    @staticmethod
+    def delete_folder_if_exists(path: str) -> None:
+        if os.path.exists(path):
+            shutil.rmtree(path)
+            print(f"Deleted folder: {path}")
+        else:
+            print(f"Folder does not exist: {path}")
+
     @staticmethod
     def resolve_path(path) -> Path | None:
         if path is None: