Additional Memory clean up for slow tests (huggingface#7436)

* update

* update

* update

Author: DN6

tests/pipelines/deepfloyd_if/test_if.py

@@ -85,6 +85,12 @@ def test_xformers_attention_forwardGenerator_pass(self):
 @slow
 @require_torch_gpu
 class IFPipelineSlowTests(unittest.TestCase):
+    def setUp(self):
+        # clean up the VRAM before each test
+        super().setUp()
+        gc.collect()
+        torch.cuda.empty_cache()
+
     def tearDown(self):
         # clean up the VRAM after each test
         super().tearDown()

tests/pipelines/deepfloyd_if/test_if_img2img.py

@@ -94,6 +94,12 @@ def test_inference_batch_single_identical(self):
 @slow
 @require_torch_gpu
 class IFImg2ImgPipelineSlowTests(unittest.TestCase):
+    def setUp(self):
+        # clean up the VRAM before each test
+        super().setUp()
+        gc.collect()
+        torch.cuda.empty_cache()
+
     def tearDown(self):
         # clean up the VRAM after each test
         super().tearDown()
@@ -109,6 +115,10 @@ def test_if_img2img(self):
         pipe.unet.set_attn_processor(AttnAddedKVProcessor())
         pipe.enable_model_cpu_offload()
 
+        torch.cuda.reset_max_memory_allocated()
+        torch.cuda.empty_cache()
+        torch.cuda.reset_peak_memory_stats()
+
         image = floats_tensor((1, 3, 64, 64), rng=random.Random(0)).to(torch_device)
         generator = torch.Generator(device="cpu").manual_seed(0)
         output = pipe(

tests/pipelines/deepfloyd_if/test_if_img2img_superresolution.py

@@ -92,6 +92,12 @@ def test_inference_batch_single_identical(self):
 @slow
 @require_torch_gpu
 class IFImg2ImgSuperResolutionPipelineSlowTests(unittest.TestCase):
+    def setUp(self):
+        # clean up the VRAM before each test
+        super().setUp()
+        gc.collect()
+        torch.cuda.empty_cache()
+
     def tearDown(self):
         # clean up the VRAM after each test
         super().tearDown()
@@ -107,6 +113,10 @@ def test_if_img2img_superresolution(self):
         pipe.unet.set_attn_processor(AttnAddedKVProcessor())
         pipe.enable_model_cpu_offload()
 
+        torch.cuda.reset_max_memory_allocated()
+        torch.cuda.empty_cache()
+        torch.cuda.reset_peak_memory_stats()
+
         generator = torch.Generator(device="cpu").manual_seed(0)
 
         original_image = floats_tensor((1, 3, 256, 256), rng=random.Random(0)).to(torch_device)

tests/pipelines/deepfloyd_if/test_if_inpainting.py

@@ -92,6 +92,12 @@ def test_inference_batch_single_identical(self):
 @slow
 @require_torch_gpu
 class IFInpaintingPipelineSlowTests(unittest.TestCase):
+    def setUp(self):
+        # clean up the VRAM before each test
+        super().setUp()
+        gc.collect()
+        torch.cuda.empty_cache()
+
     def tearDown(self):
         # clean up the VRAM after each test
         super().tearDown()
@@ -105,7 +111,6 @@ def test_if_inpainting(self):
         pipe.unet.set_attn_processor(AttnAddedKVProcessor())
         pipe.enable_model_cpu_offload()
 
-        # Super resolution test
         torch.cuda.empty_cache()
         torch.cuda.reset_max_memory_allocated()
         torch.cuda.reset_peak_memory_stats()

tests/pipelines/deepfloyd_if/test_if_inpainting_superresolution.py

@@ -94,6 +94,12 @@ def test_inference_batch_single_identical(self):
 @slow
 @require_torch_gpu
 class IFInpaintingSuperResolutionPipelineSlowTests(unittest.TestCase):
+    def setUp(self):
+        # clean up the VRAM before each test
+        super().setUp()
+        gc.collect()
+        torch.cuda.empty_cache()
+
     def tearDown(self):
         # clean up the VRAM after each test
         super().tearDown()

tests/pipelines/deepfloyd_if/test_if_superresolution.py

@@ -87,6 +87,12 @@ def test_inference_batch_single_identical(self):
 @slow
 @require_torch_gpu
 class IFSuperResolutionPipelineSlowTests(unittest.TestCase):
+    def setUp(self):
+        # clean up the VRAM before each test
+        super().setUp()
+        gc.collect()
+        torch.cuda.empty_cache()
+
     def tearDown(self):
         # clean up the VRAM after each test
         super().tearDown()

tests/pipelines/ip_adapters/test_ip_adapter_stable_diffusion.py

@@ -50,7 +50,14 @@
 class IPAdapterNightlyTestsMixin(unittest.TestCase):
     dtype = torch.float16
 
+    def setUp(self):
+        # clean up the VRAM before each test
+        super().setUp()
+        gc.collect()
+        torch.cuda.empty_cache()
+
     def tearDown(self):
+        # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()
@@ -313,7 +320,7 @@ def test_text_to_image_sdxl(self):
             feature_extractor=feature_extractor,
             torch_dtype=self.dtype,
         )
-        pipeline.to(torch_device)
+        pipeline.enable_model_cpu_offload()
         pipeline.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name="ip-adapter_sdxl.bin")
 
         inputs = self.get_dummy_inputs()
@@ -373,7 +380,7 @@ def test_image_to_image_sdxl(self):
             feature_extractor=feature_extractor,
             torch_dtype=self.dtype,
         )
-        pipeline.to(torch_device)
+        pipeline.enable_model_cpu_offload()
         pipeline.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name="ip-adapter_sdxl.bin")
 
         inputs = self.get_dummy_inputs(for_image_to_image=True)
@@ -442,7 +449,7 @@ def test_inpainting_sdxl(self):
             feature_extractor=feature_extractor,
             torch_dtype=self.dtype,
         )
-        pipeline.to(torch_device)
+        pipeline.enable_model_cpu_offload()
         pipeline.load_ip_adapter("h94/IP-Adapter", subfolder="sdxl_models", weight_name="ip-adapter_sdxl.bin")
 
         inputs = self.get_dummy_inputs(for_inpainting=True)
@@ -490,7 +497,7 @@ def test_ip_adapter_single_mask(self):
             image_encoder=image_encoder,
             torch_dtype=self.dtype,
         )
-        pipeline.to(torch_device)
+        pipeline.enable_model_cpu_offload()
         pipeline.load_ip_adapter(
             "h94/IP-Adapter", subfolder="sdxl_models", weight_name="ip-adapter-plus-face_sdxl_vit-h.safetensors"
         )
@@ -518,7 +525,7 @@ def test_ip_adapter_multiple_masks(self):
             image_encoder=image_encoder,
             torch_dtype=self.dtype,
         )
-        pipeline.to(torch_device)
+        pipeline.enable_model_cpu_offload()
         pipeline.load_ip_adapter(
             "h94/IP-Adapter", subfolder="sdxl_models", weight_name=["ip-adapter-plus-face_sdxl_vit-h.safetensors"] * 2
         )

tests/pipelines/kandinsky/test_kandinsky.py

@@ -275,6 +275,12 @@ def test_offloads(self):
 @slow
 @require_torch_gpu
 class KandinskyPipelineIntegrationTests(unittest.TestCase):
+    def setUp(self):
+        # clean up the VRAM before each test
+        super().setUp()
+        gc.collect()
+        torch.cuda.empty_cache()
+
     def tearDown(self):
         # clean up the VRAM after each test
         super().tearDown()

tests/pipelines/kandinsky/test_kandinsky_img2img.py

@@ -299,6 +299,12 @@ def test_dict_tuple_outputs_equivalent(self):
 @slow
 @require_torch_gpu
 class KandinskyImg2ImgPipelineIntegrationTests(unittest.TestCase):
+    def setUp(self):
+        # clean up the VRAM before each test
+        super().setUp()
+        gc.collect()
+        torch.cuda.empty_cache()
+
     def tearDown(self):
         # clean up the VRAM after each test
         super().tearDown()

tests/pipelines/kandinsky/test_kandinsky_inpaint.py

@@ -297,6 +297,12 @@ def test_float16_inference(self):
 @nightly
 @require_torch_gpu
 class KandinskyInpaintPipelineIntegrationTests(unittest.TestCase):
+    def setUp(self):
+        # clean up the VRAM before each test
+        super().setUp()
+        gc.collect()
+        torch.cuda.empty_cache()
+
     def tearDown(self):
         # clean up the VRAM after each test
         super().tearDown()