Merge pull request #133 from filipstrand/pre-test-cleanup

Pre test cleanup

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/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: