remove rocm from diffusers tests

huggingface · Jan 13, 2025 · 0a7a23d · 0a7a23d
1 parent eb25460
commit 0a7a23d
Showing 1 changed file with 12 additions and 6 deletions.
diff --git a/tests/onnxruntime/test_diffusion.py b/tests/onnxruntime/test_diffusion.py
@@ -281,16 +281,18 @@ def test_negative_prompt(self, model_arch: str):
         grid_parameters(
             {
                 "model_arch": SUPPORTED_ARCHITECTURES,
-                "provider": ["CUDAExecutionProvider", "ROCMExecutionProvider", "TensorrtExecutionProvider"],
+                "provider": ["CUDAExecutionProvider", "TensorrtExecutionProvider"],
             }
         )
     )
-    @pytest.mark.rocm_ep_test
     @pytest.mark.cuda_ep_test
     @pytest.mark.trt_ep_test
     @require_torch_gpu
     @require_diffusers
     def test_pipeline_on_gpu(self, test_name: str, model_arch: str, provider: str):
+        if provider == "TensorrtExecutionProvider" and model_arch != self.__class__.SUPPORTED_ARCHITECTURES[0]:
+            self.skipTest("Testing a single arch for TensorrtExecutionProvider")
+
         model_args = {"test_name": test_name, "model_arch": model_arch}
         self._setup(model_args)
 
@@ -519,16 +521,18 @@ def test_image_reproducibility(self, model_arch: str):
         grid_parameters(
             {
                 "model_arch": SUPPORTED_ARCHITECTURES,
-                "provider": ["CUDAExecutionProvider", "ROCMExecutionProvider", "TensorrtExecutionProvider"],
+                "provider": ["CUDAExecutionProvider", "TensorrtExecutionProvider"],
             }
         )
     )
-    @pytest.mark.rocm_ep_test
     @pytest.mark.cuda_ep_test
     @pytest.mark.trt_ep_test
     @require_torch_gpu
     @require_diffusers
     def test_pipeline_on_gpu(self, test_name: str, model_arch: str, provider: str):
+        if provider == "TensorrtExecutionProvider" and model_arch != self.__class__.SUPPORTED_ARCHITECTURES[0]:
+            self.skipTest("Testing a single arch for TensorrtExecutionProvider")
+
         model_args = {"test_name": test_name, "model_arch": model_arch}
         self._setup(model_args)
 
@@ -759,16 +763,18 @@ def test_image_reproducibility(self, model_arch: str):
         grid_parameters(
             {
                 "model_arch": SUPPORTED_ARCHITECTURES,
-                "provider": ["CUDAExecutionProvider", "ROCMExecutionProvider", "TensorrtExecutionProvider"],
+                "provider": ["CUDAExecutionProvider", "TensorrtExecutionProvider"],
             }
         )
     )
-    @pytest.mark.rocm_ep_test
     @pytest.mark.cuda_ep_test
     @pytest.mark.trt_ep_test
     @require_torch_gpu
     @require_diffusers
     def test_pipeline_on_gpu(self, test_name: str, model_arch: str, provider: str):
+        if provider == "TensorrtExecutionProvider" and model_arch != self.__class__.SUPPORTED_ARCHITECTURES[0]:
+            self.skipTest("Testing a single arch for TensorrtExecutionProvider")
+
         model_args = {"test_name": test_name, "model_arch": model_arch}
         self._setup(model_args)