up

Author: sayakpaul

tests/modular_pipelines/test_modular_pipelines_common.py

@@ -10,7 +10,7 @@
 from diffusers.guiders import ClassifierFreeGuidance
 from diffusers.utils import logging
 
-from ..testing_utils import backend_empty_cache, numpy_cosine_similarity_distance, torch_device
+from ..testing_utils import backend_empty_cache, numpy_cosine_similarity_distance, require_accelerator, torch_device
 
 
 class ModularPipelineTesterMixin:
@@ -197,7 +197,7 @@ def test_inference_batch_single_identical(
         max_diff = torch.abs(output_batch[0] - output[0]).max()
         assert max_diff < expected_max_diff, "Batch inference results different from single inference results"
 
-    @pytest.mark.skipif(torch_device == "cpu", reason="Test needs an accelerator.")
+    @require_accelerator
     def test_float16_inference(self, expected_max_diff=5e-2):
         pipe = self.get_pipeline()
         pipe.to(torch_device, torch.float32)
@@ -223,7 +223,7 @@ def test_float16_inference(self, expected_max_diff=5e-2):
         max_diff = numpy_cosine_similarity_distance(output.flatten(), output_fp16.flatten())
         assert max_diff < expected_max_diff, "FP16 inference is different from FP32 inference"
 
-    @pytest.mark.skipif(torch_device == "cpu", reason="Test needs an accelerator.")
+    @require_accelerator
     def test_to_device(self):
         pipe = self.get_pipeline().to("cpu")
 
@@ -246,7 +246,7 @@ def test_inference_is_not_nan_cpu(self):
         output = pipe(**self.get_dummy_inputs(), output="images")
         assert torch.isnan(output).sum() == 0, "CPU Inference returns NaN"
 
-    @pytest.mark.skipif(torch_device == "cpu", reason="Test needs an accelerator.")
+    @require_accelerator
     def test_inference_is_not_nan(self):
         pipe = self.get_pipeline().to(torch_device)
 
@@ -274,7 +274,7 @@ def test_num_images_per_prompt(self):
 
                 assert images.shape[0] == batch_size * num_images_per_prompt
 
-    @pytest.mark.skipif(torch_device == "cpu", reason="Test needs an accelerator.")
+    @require_accelerator
     def test_components_auto_cpu_offload_inference_consistent(self):
         base_pipe = self.get_pipeline().to(torch_device)