update

Author: faaany

tests/pipelines/animatediff/test_animatediff.py

@@ -1,4 +1,3 @@
-import gc
 import unittest
 
 import numpy as np
@@ -20,7 +19,7 @@
 from diffusers.models.attention import FreeNoiseTransformerBlock
 from diffusers.utils import is_xformers_available, logging
 from diffusers.utils.testing_utils import (
-    backend_empty_cache,
+    flush_memory,
     numpy_cosine_similarity_distance,
     require_accelerator,
     require_torch_accelerator,
@@ -553,14 +552,12 @@ class AnimateDiffPipelineSlowTests(unittest.TestCase):
     def setUp(self):
         # clean up the VRAM before each test
         super().setUp()
-        gc.collect()
-        backend_empty_cache(torch_device)
+        flush_memory(torch_device, gc_collect=True)
 
     def tearDown(self):
         # clean up the VRAM after each test
         super().tearDown()
-        gc.collect()
-        backend_empty_cache(torch_device)
+        flush_memory(torch_device, gc_collect=True)
 
     def test_animatediff(self):
         adapter = MotionAdapter.from_pretrained("guoyww/animatediff-motion-adapter-v1-5-2")

tests/pipelines/cogvideo/test_cogvideox_image2video.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import gc
 import inspect
 import unittest
 
@@ -24,8 +23,8 @@
 from diffusers import AutoencoderKLCogVideoX, CogVideoXImageToVideoPipeline, CogVideoXTransformer3DModel, DDIMScheduler
 from diffusers.utils import load_image
 from diffusers.utils.testing_utils import (
-    backend_empty_cache,
     enable_full_determinism,
+    flush_memory,
     numpy_cosine_similarity_distance,
     require_torch_accelerator,
     slow,
@@ -351,13 +350,11 @@ class CogVideoXImageToVideoPipelineIntegrationTests(unittest.TestCase):
 
     def setUp(self):
         super().setUp()
-        gc.collect()
-        backend_empty_cache(torch_device)
+        flush_memory(torch_device, gc_collect=True)
 
     def tearDown(self):
         super().tearDown()
-        gc.collect()
-        backend_empty_cache(torch_device)
+        flush_memory(torch_device, gc_collect=True)
 
     def test_cogvideox(self):
         generator = torch.Generator("cpu").manual_seed(0)

tests/pipelines/controlnet/test_controlnet.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import gc
 import tempfile
 import traceback
 import unittest
@@ -34,8 +33,8 @@
 from diffusers.pipelines.controlnet.pipeline_controlnet import MultiControlNetModel
 from diffusers.utils.import_utils import is_xformers_available
 from diffusers.utils.testing_utils import (
-    backend_empty_cache,
     enable_full_determinism,
+    flush_memory,
     get_python_version,
     is_torch_compile,
     load_image,
@@ -704,13 +703,11 @@ def test_save_pretrained_raise_not_implemented_exception(self):
 class ControlNetPipelineSlowTests(unittest.TestCase):
     def setUp(self):
         super().setUp()
-        gc.collect()
-        backend_empty_cache(torch_device)
+        flush_memory(torch_device, gc_collect=True)
 
     def tearDown(self):
         super().tearDown()
-        gc.collect()
-        backend_empty_cache(torch_device)
+        flush_memory(torch_device, gc_collect=True)
 
     def test_canny(self):
         controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny")
@@ -929,14 +926,7 @@ def test_seg(self):
         assert np.abs(expected_image - image).max() < 8e-2
 
     def test_sequential_cpu_offloading(self):
-        if torch_device == "cuda":
-            torch.cuda.empty_cache()
-            torch.cuda.reset_max_memory_allocated()
-            torch.cuda.reset_peak_memory_stats()
-        elif torch_device == "xpu":
-            torch.xpu.empty_cache()
-            torch.xpu.reset_max_memory_allocated()
-            torch.xpu.reset_peak_memory_stats()
+        flush_memory(torch_device, gc_collect=True, reset_mem_stats=True)
 
         controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-seg")
 
@@ -1077,13 +1067,11 @@ def test_v11_shuffle_global_pool_conditions(self):
 class StableDiffusionMultiControlNetPipelineSlowTests(unittest.TestCase):
     def setUp(self):
         super().setUp()
-        gc.collect()
-        backend_empty_cache(torch_device)
+        flush_memory(torch_device, gc_collect=True)
 
     def tearDown(self):
         super().tearDown()
-        gc.collect()
-        backend_empty_cache(torch_device)
+        flush_memory(torch_device, gc_collect=True)
 
     def test_pose_and_canny(self):
         controlnet_canny = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny")

tests/pipelines/controlnet/test_controlnet_sdxl.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 
 import copy
-import gc
 import unittest
 
 import numpy as np
@@ -35,8 +34,8 @@
 from diffusers.pipelines.controlnet.pipeline_controlnet import MultiControlNetModel
 from diffusers.utils.import_utils import is_xformers_available
 from diffusers.utils.testing_utils import (
-    backend_empty_cache,
     enable_full_determinism,
+    flush_memory,
     load_image,
     require_torch_accelerator,
     slow,
@@ -894,13 +893,11 @@ def test_negative_conditions(self):
 class ControlNetSDXLPipelineSlowTests(unittest.TestCase):
     def setUp(self):
         super().setUp()
-        gc.collect()
-        backend_empty_cache(torch_device)
+        flush_memory(torch_device, gc_collect=True)
 
     def tearDown(self):
         super().tearDown()
-        gc.collect()
-        backend_empty_cache(torch_device)
+        flush_memory(torch_device, gc_collect=True)
 
     def test_canny(self):
         controlnet = ControlNetModel.from_pretrained("diffusers/controlnet-canny-sdxl-1.0")

tests/pipelines/controlnet_hunyuandit/test_controlnet_hunyuandit.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import gc
 import unittest
 
 import numpy as np
@@ -29,8 +28,8 @@
 from diffusers.models import HunyuanDiT2DControlNetModel, HunyuanDiT2DMultiControlNetModel
 from diffusers.utils import load_image
 from diffusers.utils.testing_utils import (
-    backend_empty_cache,
     enable_full_determinism,
+    flush_memory,
     require_torch_accelerator,
     slow,
     torch_device,
@@ -185,13 +184,11 @@ class HunyuanDiTControlNetPipelineSlowTests(unittest.TestCase):
 
     def setUp(self):
         super().setUp()
-        gc.collect()
-        backend_empty_cache(torch_device)
+        flush_memory(torch_device, gc_collect=True)
 
     def tearDown(self):
         super().tearDown()
-        gc.collect()
-        backend_empty_cache(torch_device)
+        flush_memory(torch_device, gc_collect=True)
 
     def test_canny(self):
         controlnet = HunyuanDiT2DControlNetModel.from_pretrained(

tests/pipelines/controlnet_xs/test_controlnetxs.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import gc
 import traceback
 import unittest
 
@@ -34,8 +33,8 @@
 )
 from diffusers.utils.import_utils import is_xformers_available
 from diffusers.utils.testing_utils import (
-    backend_empty_cache,
     enable_full_determinism,
+    flush_memory,
     is_torch_compile,
     load_image,
     load_numpy,
@@ -339,8 +338,7 @@ def test_to_device(self):
 class ControlNetXSPipelineSlowTests(unittest.TestCase):
     def tearDown(self):
         super().tearDown()
-        gc.collect()
-        backend_empty_cache(torch_device)
+        flush_memory(torch_device, gc_collect=True)
 
     def test_canny(self):
         controlnet = ControlNetXSAdapter.from_pretrained(

tests/pipelines/controlnet_xs/test_controlnetxs_sdxl.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import gc
 import unittest
 
 import numpy as np
@@ -32,8 +31,8 @@
 )
 from diffusers.utils.import_utils import is_xformers_available
 from diffusers.utils.testing_utils import (
-    backend_empty_cache,
     enable_full_determinism,
+    flush_memory,
     load_image,
     require_torch_accelerator,
     slow,
@@ -380,8 +379,7 @@ def test_multi_vae(self):
 class StableDiffusionXLControlNetXSPipelineSlowTests(unittest.TestCase):
     def tearDown(self):
         super().tearDown()
-        gc.collect()
-        backend_empty_cache(torch_device)
+        flush_memory(torch_device, gc_collect=True)
 
     def test_canny(self):
         controlnet = ControlNetXSAdapter.from_pretrained(

tests/pipelines/deepfloyd_if/test_if_img2img.py

@@ -24,6 +24,7 @@
 from diffusers.utils.import_utils import is_xformers_available
 from diffusers.utils.testing_utils import (
     floats_tensor,
+    flush_memory,
     load_numpy,
     require_accelerator,
     require_torch_accelerator,
@@ -125,14 +126,7 @@ def test_if_img2img(self):
         pipe.unet.set_attn_processor(AttnAddedKVProcessor())
         pipe.enable_model_cpu_offload(device=torch_device)
 
-        if torch_device == "cuda":
-            torch.cuda.reset_max_memory_allocated()
-            torch.cuda.empty_cache()
-            torch.cuda.reset_peak_memory_stats()
-        elif torch_device == "xpu":
-            torch.xpu.reset_max_memory_allocated()
-            torch.xpu.empty_cache()
-            torch.xpu.reset_peak_memory_stats()
+        flush_memory(torch_device, reset_mem_stats=True)
 
         image = floats_tensor((1, 3, 64, 64), rng=random.Random(0)).to(torch_device)
         generator = torch.Generator(device="cpu").manual_seed(0)

tests/pipelines/deepfloyd_if/test_if_img2img_superresolution.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import gc
 import random
 import unittest
 
@@ -23,8 +22,8 @@
 from diffusers.models.attention_processor import AttnAddedKVProcessor
 from diffusers.utils.import_utils import is_xformers_available
 from diffusers.utils.testing_utils import (
-    backend_empty_cache,
     floats_tensor,
+    flush_memory,
     load_numpy,
     require_accelerator,
     require_torch_accelerator,
@@ -105,14 +104,12 @@ class IFImg2ImgSuperResolutionPipelineSlowTests(unittest.TestCase):
     def setUp(self):
         # clean up the VRAM before each test
         super().setUp()
-        gc.collect()
-        backend_empty_cache(torch_device)
+        flush_memory(torch_device, gc_collect=True)
 
     def tearDown(self):
         # clean up the VRAM after each test
         super().tearDown()
-        gc.collect()
-        torch.cuda.empty_cache()
+        flush_memory(torch_device, gc_collect=True)
 
     def test_if_img2img_superresolution(self):
         pipe = IFImg2ImgSuperResolutionPipeline.from_pretrained(
@@ -123,14 +120,7 @@ def test_if_img2img_superresolution(self):
         pipe.unet.set_attn_processor(AttnAddedKVProcessor())
         pipe.enable_model_cpu_offload(device=torch_device)
 
-        if torch_device == "cuda":
-            torch.cuda.reset_max_memory_allocated()
-            torch.cuda.empty_cache()
-            torch.cuda.reset_peak_memory_stats()
-        elif torch_device == "xpu":
-            torch.xpu.reset_max_memory_allocated()
-            torch.xpu.empty_cache()
-            torch.xpu.reset_peak_memory_stats()
+        flush_memory(torch_device, reset_mem_stats=True)
 
         generator = torch.Generator(device="cpu").manual_seed(0)
 

tests/pipelines/i2vgen_xl/test_i2vgenxl.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import gc
 import random
 import unittest
 
@@ -36,9 +35,9 @@
 from diffusers.models.unets import I2VGenXLUNet
 from diffusers.utils import is_xformers_available, load_image
 from diffusers.utils.testing_utils import (
-    backend_empty_cache,
     enable_full_determinism,
     floats_tensor,
+    flush_memory,
     numpy_cosine_similarity_distance,
     require_torch_accelerator,
     skip_mps,
@@ -232,14 +231,12 @@ class I2VGenXLPipelineSlowTests(unittest.TestCase):
     def setUp(self):
         # clean up the VRAM before each test
         super().setUp()
-        gc.collect()
-        backend_empty_cache(torch_device)
+        flush_memory(torch_device, gc_collect=True)
 
     def tearDown(self):
         # clean up the VRAM after each test
         super().tearDown()
-        gc.collect()
-        backend_empty_cache(torch_device)
+        flush_memory(torch_device, gc_collect=True)
 
     def test_i2vgen_xl(self):
         pipe = I2VGenXLPipeline.from_pretrained("ali-vilab/i2vgen-xl", torch_dtype=torch.float16, variant="fp16")