Skip to content

Commit 1967c62

Browse files
yao-matrixyao-matrix
committed
device agnostic APIs
Signed-off-by: YAO Matrix <[email protected]>
1 parent 2e11423 commit 1967c62

23 files changed

+63
-46
lines changed

‎tests/models/autoencoders/test_models_consistency_decoder_vae.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
from diffusers import ConsistencyDecoderVAE, StableDiffusionPipeline
2323
from diffusers.utils.testing_utils import (
24+
backend_empty_cache,
2425
enable_full_determinism,
2526
load_image,
2627
slow,
@@ -162,13 +163,13 @@ def setUp(self):
162163
# clean up the VRAM before each test
163164
super().setUp()
164165
gc.collect()
165-
torch.cuda.empty_cache()
166+
backend_empty_cache(torch_device)
166167

167168
def tearDown(self):
168169
# clean up the VRAM after each test
169170
super().tearDown()
170171
gc.collect()
171-
torch.cuda.empty_cache()
172+
backend_empty_cache(torch_device)
172173

173174
@torch.no_grad()
174175
def test_encode_decode(self):

‎tests/models/unets/test_models_unet_2d.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
from diffusers import UNet2DModel
2323
from diffusers.utils import logging
2424
from diffusers.utils.testing_utils import (
25+
backend_empty_cache,
2526
enable_full_determinism,
2627
floats_tensor,
2728
require_torch_accelerator,
@@ -229,7 +230,7 @@ def test_from_pretrained_accelerate_wont_change_results(self):
229230

230231
# two models don't need to stay in the device at the same time
231232
del model_accelerate
232-
torch.cuda.empty_cache()
233+
backend_empty_cache(torch_device)
233234
gc.collect()
234235

235236
model_normal_load, _ = UNet2DModel.from_pretrained(

‎tests/pipelines/allegro/test_allegro.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424

2525
from diffusers import AllegroPipeline, AllegroTransformer3DModel, AutoencoderKLAllegro, DDIMScheduler
2626
from diffusers.utils.testing_utils import (
27+
backend_empty_cache,
2728
enable_full_determinism,
2829
numpy_cosine_similarity_distance,
2930
require_hf_hub_version_greater,
@@ -341,12 +342,12 @@ class AllegroPipelineIntegrationTests(unittest.TestCase):
341342
def setUp(self):
342343
super().setUp()
343344
gc.collect()
344-
torch.cuda.empty_cache()
345+
backend_empty_cache(torch_device)
345346

346347
def tearDown(self):
347348
super().tearDown()
348349
gc.collect()
349-
torch.cuda.empty_cache()
350+
backend_empty_cache(torch_device)
350351

351352
def test_allegro(self):
352353
generator = torch.Generator("cpu").manual_seed(0)

‎tests/pipelines/audioldm/test_audioldm.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@
3737
UNet2DConditionModel,
3838
)
3939
from diffusers.utils import is_xformers_available
40-
from diffusers.utils.testing_utils import enable_full_determinism, nightly, torch_device
40+
from diffusers.utils.testing_utils import backend_empty_cache, enable_full_determinism, nightly, torch_device
4141

4242
from ..pipeline_params import TEXT_TO_AUDIO_BATCH_PARAMS, TEXT_TO_AUDIO_PARAMS
4343
from ..test_pipelines_common import PipelineTesterMixin
@@ -378,12 +378,12 @@ class AudioLDMPipelineSlowTests(unittest.TestCase):
378378
def setUp(self):
379379
super().setUp()
380380
gc.collect()
381-
torch.cuda.empty_cache()
381+
backend_empty_cache(torch_device)
382382

383383
def tearDown(self):
384384
super().tearDown()
385385
gc.collect()
386-
torch.cuda.empty_cache()
386+
backend_empty_cache(torch_device)
387387

388388
def get_inputs(self, device, generator_device="cpu", dtype=torch.float32, seed=0):
389389
generator = torch.Generator(device=generator_device).manual_seed(seed)
@@ -423,12 +423,12 @@ class AudioLDMPipelineNightlyTests(unittest.TestCase):
423423
def setUp(self):
424424
super().setUp()
425425
gc.collect()
426-
torch.cuda.empty_cache()
426+
backend_empty_cache(torch_device)
427427

428428
def tearDown(self):
429429
super().tearDown()
430430
gc.collect()
431-
torch.cuda.empty_cache()
431+
backend_empty_cache(torch_device)
432432

433433
def get_inputs(self, device, generator_device="cpu", dtype=torch.float32, seed=0):
434434
generator = torch.Generator(device=generator_device).manual_seed(seed)

‎tests/pipelines/audioldm2/test_audioldm2.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@
4545
LMSDiscreteScheduler,
4646
PNDMScheduler,
4747
)
48-
from diffusers.utils.testing_utils import enable_full_determinism, is_torch_version, nightly, torch_device
48+
from diffusers.utils.testing_utils import backend_empty_cache, enable_full_determinism, is_torch_version, nightly, torch_device
4949

5050
from ..pipeline_params import TEXT_TO_AUDIO_BATCH_PARAMS, TEXT_TO_AUDIO_PARAMS
5151
from ..test_pipelines_common import PipelineTesterMixin
@@ -540,12 +540,12 @@ class AudioLDM2PipelineSlowTests(unittest.TestCase):
540540
def setUp(self):
541541
super().setUp()
542542
gc.collect()
543-
torch.cuda.empty_cache()
543+
backend_empty_cache(torch_device)
544544

545545
def tearDown(self):
546546
super().tearDown()
547547
gc.collect()
548-
torch.cuda.empty_cache()
548+
backend_empty_cache(torch_device)
549549

550550
def get_inputs(self, device, generator_device="cpu", dtype=torch.float32, seed=0):
551551
generator = torch.Generator(device=generator_device).manual_seed(seed)

‎tests/pipelines/cogvideo/test_cogvideox.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222

2323
from diffusers import AutoencoderKLCogVideoX, CogVideoXPipeline, CogVideoXTransformer3DModel, DDIMScheduler
2424
from diffusers.utils.testing_utils import (
25+
backend_empty_cache,
2526
enable_full_determinism,
2627
numpy_cosine_similarity_distance,
2728
require_torch_accelerator,
@@ -334,12 +335,12 @@ class CogVideoXPipelineIntegrationTests(unittest.TestCase):
334335
def setUp(self):
335336
super().setUp()
336337
gc.collect()
337-
torch.cuda.empty_cache()
338+
backend_empty_cache(torch_device)
338339

339340
def tearDown(self):
340341
super().tearDown()
341342
gc.collect()
342-
torch.cuda.empty_cache()
343+
backend_empty_cache(torch_device)
343344

344345
def test_cogvideox(self):
345346
generator = torch.Generator("cpu").manual_seed(0)

‎tests/pipelines/cogview3/test_cogview3plus.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222

2323
from diffusers import AutoencoderKL, CogVideoXDDIMScheduler, CogView3PlusPipeline, CogView3PlusTransformer2DModel
2424
from diffusers.utils.testing_utils import (
25+
backend_empty_cache,
2526
enable_full_determinism,
2627
numpy_cosine_similarity_distance,
2728
require_torch_accelerator,
@@ -244,12 +245,12 @@ class CogView3PlusPipelineIntegrationTests(unittest.TestCase):
244245
def setUp(self):
245246
super().setUp()
246247
gc.collect()
247-
torch.cuda.empty_cache()
248+
backend_empty_cache(torch_device)
248249

249250
def tearDown(self):
250251
super().tearDown()
251252
gc.collect()
252-
torch.cuda.empty_cache()
253+
backend_empty_cache(torch_device)
253254

254255
def test_cogview3plus(self):
255256
generator = torch.Generator("cpu").manual_seed(0)

‎tests/pipelines/controlnet/test_controlnet_img2img.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
from diffusers.utils import load_image
3737
from diffusers.utils.import_utils import is_xformers_available
3838
from diffusers.utils.testing_utils import (
39+
backend_empty_cache,
3940
enable_full_determinism,
4041
floats_tensor,
4142
load_numpy,
@@ -412,12 +413,12 @@ class ControlNetImg2ImgPipelineSlowTests(unittest.TestCase):
412413
def setUp(self):
413414
super().setUp()
414415
gc.collect()
415-
torch.cuda.empty_cache()
416+
backend_empty_cache(torch_device)
416417

417418
def tearDown(self):
418419
super().tearDown()
419420
gc.collect()
420-
torch.cuda.empty_cache()
421+
backend_empty_cache(torch_device)
421422

422423
def test_canny(self):
423424
controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny")

‎tests/pipelines/controlnet/test_controlnet_inpaint.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
from diffusers.utils import load_image
3737
from diffusers.utils.import_utils import is_xformers_available
3838
from diffusers.utils.testing_utils import (
39+
backend_empty_cache,
3940
enable_full_determinism,
4041
floats_tensor,
4142
load_numpy,
@@ -464,12 +465,12 @@ class ControlNetInpaintPipelineSlowTests(unittest.TestCase):
464465
def setUp(self):
465466
super().setUp()
466467
gc.collect()
467-
torch.cuda.empty_cache()
468+
backend_empty_cache(torch_device)
468469

469470
def tearDown(self):
470471
super().tearDown()
471472
gc.collect()
472-
torch.cuda.empty_cache()
473+
backend_empty_cache(torch_device)
473474

474475
def test_canny(self):
475476
controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny")

‎tests/pipelines/controlnet_sd3/test_controlnet_sd3.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -221,7 +221,7 @@ def test_xformers_attention_forwardGenerator_pass(self):
221221

222222
@slow
223223
@require_big_accelerator
224-
@pytest.mark.big_gpu_with_torch_cuda
224+
@pytest.mark.big_accelerator
225225
class StableDiffusion3ControlNetPipelineSlowTests(unittest.TestCase):
226226
pipeline_class = StableDiffusion3ControlNetPipeline
227227

0 commit comments

Comments
 (0)