From d80b6e520fc29bedd82e3e43fae8aff88e1d2d03 Mon Sep 17 00:00:00 2001 From: filipstrand Date: Sat, 1 Mar 2025 21:32:14 +0100 Subject: [PATCH 1/3] Add option to override saved images --- src/mflux/post_processing/generated_image.py | 9 ++++++-- src/mflux/post_processing/image_util.py | 22 +++++++++++--------- 2 files changed, 19 insertions(+), 12 deletions(-) diff --git a/src/mflux/post_processing/generated_image.py b/src/mflux/post_processing/generated_image.py index caee027..33f8859 100644 --- a/src/mflux/post_processing/generated_image.py +++ b/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 diff --git a/src/mflux/post_processing/image_util.py b/src/mflux/post_processing/image_util.py index a21da15..eae15c1 100644 --- a/src/mflux/post_processing/image_util.py +++ b/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 From 09430350f1d9baefa081af1ee12bec85ceab9248 Mon Sep 17 00:00:00 2001 From: filipstrand Date: Sat, 1 Mar 2025 21:37:03 +0100 Subject: [PATCH 2/3] Override previously saved test images --- .../helpers/image_generation_controlnet_test_helper.py | 2 +- tests/image_generation/helpers/image_generation_test_helper.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/image_generation/helpers/image_generation_controlnet_test_helper.py b/tests/image_generation/helpers/image_generation_controlnet_test_helper.py index f89eb1f..d132fd1 100644 --- a/tests/image_generation/helpers/image_generation_controlnet_test_helper.py +++ b/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( diff --git a/tests/image_generation/helpers/image_generation_test_helper.py b/tests/image_generation/helpers/image_generation_test_helper.py index 0ef548c..6718dc1 100644 --- a/tests/image_generation/helpers/image_generation_test_helper.py +++ b/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( From 06f336c71465928798100cf155494e7e40e62ded Mon Sep 17 00:00:00 2001 From: filipstrand Date: Sat, 1 Mar 2025 22:10:10 +0100 Subject: [PATCH 3/3] Pre test cleanup --- tests/dreambooth/test_resume_training.py | 11 +++++++++++ tests/dreambooth/test_train_and_load_weights.py | 12 ++++++++++++ tests/model_saving/test_model_saving.py | 12 ++++++++++++ tests/model_saving/test_model_saving_lora.py | 12 ++++++++++++ 4 files changed, 47 insertions(+) diff --git a/tests/dreambooth/test_resume_training.py b/tests/dreambooth/test_resume_training.py index 423119b..1d3c4c8 100644 --- a/tests/dreambooth/test_resume_training.py +++ b/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.") diff --git a/tests/dreambooth/test_train_and_load_weights.py b/tests/dreambooth/test_train_and_load_weights.py index 58163e9..bb6456e 100644 --- a/tests/dreambooth/test_train_and_load_weights.py +++ b/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.") diff --git a/tests/model_saving/test_model_saving.py b/tests/model_saving/test_model_saving.py index 82a0e7a..96cc696 100644 --- a/tests/model_saving/test_model_saving.py +++ b/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}") diff --git a/tests/model_saving/test_model_saving_lora.py b/tests/model_saving/test_model_saving_lora.py index acf3a88..88bf8ab 100644 --- a/tests/model_saving/test_model_saving_lora.py +++ b/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: