diff --git a/tests/others/test_cli_commands.py b/tests/others/test_cli_commands.py index 5b27eaf372f4..9b731416df79 100644 --- a/tests/others/test_cli_commands.py +++ b/tests/others/test_cli_commands.py @@ -49,6 +49,7 @@ from diffusers.utils.testing_utils import ( require_accelerator, require_kernels_version_greater_or_equal, + require_torch_accelerator, require_torch_gpu, ) @@ -201,7 +202,7 @@ def _parse_run_argv(self, extra_argv: list[str]) -> Namespace: ] ) - @require_torch_gpu + @require_torch_accelerator def test_group_offload_arg(self): from diffusers.hooks.group_offloading import _is_group_offload_enabled diff --git a/tests/pipelines/dreamlite/test_pipeline_dreamlite.py b/tests/pipelines/dreamlite/test_pipeline_dreamlite.py index 51a564f7caed..b186589d014b 100644 --- a/tests/pipelines/dreamlite/test_pipeline_dreamlite.py +++ b/tests/pipelines/dreamlite/test_pipeline_dreamlite.py @@ -37,9 +37,10 @@ from diffusers import DreamLitePipeline, DreamLiteUNetModel from ...testing_utils import ( + backend_empty_cache, enable_full_determinism, nightly, - require_torch_gpu, + require_torch_accelerator, torch_device, ) from ..testing_utils import MemoryTesterMixin, PipelineTesterMixin @@ -221,7 +222,7 @@ class TestDreamLitePipelineMemory(DreamLitePipelineTesterConfig, MemoryTesterMix @nightly -@require_torch_gpu +@require_torch_accelerator class TestDreamLitePipelineIntegration: """End-to-end test against the real DreamLite-base checkpoint on the Hub. @@ -236,10 +237,10 @@ class TestDreamLitePipelineIntegration: @pytest.fixture(autouse=True) def cleanup(self): gc.collect() - torch.cuda.empty_cache() + backend_empty_cache(torch_device) yield gc.collect() - torch.cuda.empty_cache() + backend_empty_cache(torch_device) def _from_pretrained_kwargs(self): local = os.getenv("DREAMLITE_BASE_PATH") @@ -249,7 +250,7 @@ def _from_pretrained_kwargs(self): def test_dreamlite_t2i_real_checkpoint(self): pipe = DreamLitePipeline.from_pretrained(**self._from_pretrained_kwargs(), torch_dtype=torch.bfloat16).to( - "cuda" + torch_device ) out = pipe( prompt="a dog running on the grass", @@ -266,7 +267,7 @@ def test_dreamlite_t2i_real_checkpoint(self): def test_dreamlite_i2i_real_checkpoint(self): pipe = DreamLitePipeline.from_pretrained(**self._from_pretrained_kwargs(), torch_dtype=torch.bfloat16).to( - "cuda" + torch_device ) src = Image.fromarray((np.random.RandomState(0).rand(1024, 1024, 3) * 255).astype(np.uint8)) diff --git a/tests/pipelines/dreamlite/test_pipeline_dreamlite_mobile.py b/tests/pipelines/dreamlite/test_pipeline_dreamlite_mobile.py index 268f2ca6ede8..de09b8f745bd 100644 --- a/tests/pipelines/dreamlite/test_pipeline_dreamlite_mobile.py +++ b/tests/pipelines/dreamlite/test_pipeline_dreamlite_mobile.py @@ -31,9 +31,10 @@ from diffusers import DreamLiteMobilePipeline from ...testing_utils import ( + backend_empty_cache, enable_full_determinism, nightly, - require_torch_gpu, + require_torch_accelerator, torch_device, ) from ..testing_utils import MemoryTesterMixin, PipelineTesterMixin @@ -153,7 +154,7 @@ class TestDreamLiteMobilePipelineMemory(DreamLiteMobilePipelineTesterConfig, Mem @nightly -@require_torch_gpu +@require_torch_accelerator class TestDreamLiteMobilePipelineIntegration: """End-to-end test against the real DreamLite-mobile checkpoint on the Hub. @@ -168,10 +169,10 @@ class TestDreamLiteMobilePipelineIntegration: @pytest.fixture(autouse=True) def cleanup(self): gc.collect() - torch.cuda.empty_cache() + backend_empty_cache(torch_device) yield gc.collect() - torch.cuda.empty_cache() + backend_empty_cache(torch_device) def _from_pretrained_kwargs(self): local = os.getenv("DREAMLITE_MOBILE_PATH") @@ -182,7 +183,7 @@ def _from_pretrained_kwargs(self): def test_mobile_t2i_real_checkpoint(self): pipe = DreamLiteMobilePipeline.from_pretrained( **self._from_pretrained_kwargs(), torch_dtype=torch.bfloat16 - ).to("cuda") + ).to(torch_device) out = pipe( prompt="a dog running on the grass", num_inference_steps=4, @@ -198,7 +199,7 @@ def test_mobile_t2i_real_checkpoint(self): def test_mobile_i2i_real_checkpoint(self): pipe = DreamLiteMobilePipeline.from_pretrained( **self._from_pretrained_kwargs(), torch_dtype=torch.bfloat16 - ).to("cuda") + ).to(torch_device) src = Image.fromarray((np.random.RandomState(0).rand(1024, 1024, 3) * 255).astype(np.uint8)) out = pipe(