with gc.collect

Author: faaany

tests/pipelines/animatediff/test_animatediff.py

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