[ONNX] Fix flaky tests (#1593)

* [ONNX] Fix flaky tests

* revert

tests/pipelines/stable_diffusion/test_onnx_stable_diffusion.py

@@ -58,6 +58,8 @@ class OnnxStableDiffusionPipelineIntegrationTests(unittest.TestCase):
         sd_pipe = OnnxStableDiffusionPipeline.from_pretrained(
             "CompVis/stable-diffusion-v1-4",
             revision="onnx",
+            safety_checker=None,
+            feature_extractor=None,
             provider=self.gpu_provider,
             sess_options=self.gpu_options,
         )
@@ -82,6 +84,8 @@ class OnnxStableDiffusionPipelineIntegrationTests(unittest.TestCase):
             "runwayml/stable-diffusion-v1-5",
             revision="onnx",
             scheduler=ddim_scheduler,
+            safety_checker=None,
+            feature_extractor=None,
             provider=self.gpu_provider,
             sess_options=self.gpu_options,
         )
@@ -105,6 +109,8 @@ class OnnxStableDiffusionPipelineIntegrationTests(unittest.TestCase):
             "runwayml/stable-diffusion-v1-5",
             revision="onnx",
             scheduler=lms_scheduler,
+            safety_checker=None,
+            feature_extractor=None,
             provider=self.gpu_provider,
             sess_options=self.gpu_options,
         )
@@ -147,6 +153,8 @@ class OnnxStableDiffusionPipelineIntegrationTests(unittest.TestCase):
         pipe = OnnxStableDiffusionPipeline.from_pretrained(
             "runwayml/stable-diffusion-v1-5",
             revision="onnx",
+            safety_checker=None,
+            feature_extractor=None,
             provider=self.gpu_provider,
             sess_options=self.gpu_options,
         )
@@ -170,9 +178,10 @@ class OnnxStableDiffusionPipelineIntegrationTests(unittest.TestCase):
         pipe = OnnxStableDiffusionPipeline.from_pretrained(
             "runwayml/stable-diffusion-v1-5",
             revision="onnx",
+            safety_checker=None,
+            feature_extractor=None,
             provider=self.gpu_provider,
             sess_options=self.gpu_options,
-            safety_checker=None,
         )
         assert isinstance(pipe, OnnxStableDiffusionPipeline)
         assert pipe.safety_checker is None

tests/pipelines/stable_diffusion/test_onnx_stable_diffusion_img2img.py

@@ -62,6 +62,8 @@ class OnnxStableDiffusionImg2ImgPipelineIntegrationTests(unittest.TestCase):
         pipe = OnnxStableDiffusionImg2ImgPipeline.from_pretrained(
             "CompVis/stable-diffusion-v1-4",
             revision="onnx",
+            safety_checker=None,
+            feature_extractor=None,
             provider=self.gpu_provider,
             sess_options=self.gpu_options,
         )
@@ -100,6 +102,8 @@ class OnnxStableDiffusionImg2ImgPipelineIntegrationTests(unittest.TestCase):
             "runwayml/stable-diffusion-v1-5",
             revision="onnx",
             scheduler=lms_scheduler,
+            safety_checker=None,
+            feature_extractor=None,
             provider=self.gpu_provider,
             sess_options=self.gpu_options,
         )
@@ -113,7 +117,7 @@ class OnnxStableDiffusionImg2ImgPipelineIntegrationTests(unittest.TestCase):
             image=init_image,
             strength=0.75,
             guidance_scale=7.5,
-            num_inference_steps=10,
+            num_inference_steps=20,
             generator=generator,
             output_type="np",
         )
@@ -121,6 +125,6 @@ class OnnxStableDiffusionImg2ImgPipelineIntegrationTests(unittest.TestCase):
         image_slice = images[0, 255:258, 383:386, -1]
 
         assert images.shape == (1, 512, 768, 3)
-        expected_slice = np.array([0.7950, 0.7923, 0.7903, 0.5516, 0.5501, 0.5476, 0.4965, 0.4933, 0.4910])
+        expected_slice = np.array([0.8043, 0.926, 0.9581, 0.8119, 0.8954, 0.913, 0.7209, 0.7463, 0.7431])
         # TODO: lower the tolerance after finding the cause of onnxruntime reproducibility issues
         assert np.abs(image_slice.flatten() - expected_slice).max() < 2e-2

tests/pipelines/stable_diffusion/test_onnx_stable_diffusion_inpaint.py

@@ -64,6 +64,8 @@ class OnnxStableDiffusionInpaintPipelineIntegrationTests(unittest.TestCase):
         pipe = OnnxStableDiffusionInpaintPipeline.from_pretrained(
             "runwayml/stable-diffusion-inpainting",
             revision="onnx",
+            safety_checker=None,
+            feature_extractor=None,
             provider=self.gpu_provider,
             sess_options=self.gpu_options,
         )
@@ -104,6 +106,8 @@ class OnnxStableDiffusionInpaintPipelineIntegrationTests(unittest.TestCase):
             "runwayml/stable-diffusion-inpainting",
             revision="onnx",
             scheduler=lms_scheduler,
+            safety_checker=None,
+            feature_extractor=None,
             provider=self.gpu_provider,
             sess_options=self.gpu_options,
         )
@@ -117,7 +121,7 @@ class OnnxStableDiffusionInpaintPipelineIntegrationTests(unittest.TestCase):
             image=init_image,
             mask_image=mask_image,
             guidance_scale=7.5,
-            num_inference_steps=10,
+            num_inference_steps=20,
             generator=generator,
             output_type="np",
         )
@@ -125,5 +129,5 @@ class OnnxStableDiffusionInpaintPipelineIntegrationTests(unittest.TestCase):
         image_slice = images[0, 255:258, 255:258, -1]
 
         assert images.shape == (1, 512, 512, 3)
-        expected_slice = np.array([0.2520, 0.2743, 0.2643, 0.2641, 0.2517, 0.2650, 0.2498, 0.2688, 0.2529])
+        expected_slice = np.array([0.0086, 0.0077, 0.0083, 0.0093, 0.0107, 0.0139, 0.0094, 0.0097, 0.0125])
         assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-3

tests/pipelines/stable_diffusion/test_onnx_stable_diffusion_inpaint_legacy.py

@@ -70,6 +70,8 @@ class StableDiffusionOnnxInpaintLegacyPipelineIntegrationTests(unittest.TestCase
         pipe = OnnxStableDiffusionInpaintPipelineLegacy.from_pretrained(
             "CompVis/stable-diffusion-v1-4",
             revision="onnx",
+            safety_checker=None,
+            feature_extractor=None,
             provider=self.gpu_provider,
             sess_options=self.gpu_options,
         )