[ONNX] Fix flaky tests (huggingface#1593)

* [ONNX] Fix flaky tests * revert
WAUthethird · Dec 7, 2022 · eb1abee · eb1abee
1 parent 5e03692
commit eb1abee
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 5 deletions.
diff --git a/tests/pipelines/stable_diffusion/test_onnx_stable_diffusion.py b/tests/pipelines/stable_diffusion/test_onnx_stable_diffusion.py
@@ -58,6 +58,8 @@ def test_inference_default_pndm(self):
         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 @@ def test_inference_ddim(self):
             "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 @@ def test_inference_k_lms(self):
             "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 @@ def test_callback_fn(step: int, timestep: int, latents: np.ndarray) -> None:
         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 @@ def test_stable_diffusion_no_safety_checker(self):
         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

diff --git a/tests/pipelines/stable_diffusion/test_onnx_stable_diffusion_img2img.py b/tests/pipelines/stable_diffusion/test_onnx_stable_diffusion_img2img.py
@@ -62,6 +62,8 @@ def test_inference_default_pndm(self):
         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 @@ def test_inference_k_lms(self):
             "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,14 +117,14 @@ def test_inference_k_lms(self):
             image=init_image,
             strength=0.75,
             guidance_scale=7.5,
-            num_inference_steps=10,
+            num_inference_steps=20,
             generator=generator,
             output_type="np",
         )
         images = output.images
         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
diff --git a/tests/pipelines/stable_diffusion/test_onnx_stable_diffusion_inpaint.py b/tests/pipelines/stable_diffusion/test_onnx_stable_diffusion_inpaint.py
@@ -64,6 +64,8 @@ def test_inference_default_pndm(self):
         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 @@ def test_inference_k_lms(self):
             "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,13 +121,13 @@ def test_inference_k_lms(self):
             image=init_image,
             mask_image=mask_image,
             guidance_scale=7.5,
-            num_inference_steps=10,
+            num_inference_steps=20,
             generator=generator,
             output_type="np",
         )
         images = output.images
         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
diff --git a/tests/pipelines/stable_diffusion/test_onnx_stable_diffusion_inpaint_legacy.py b/tests/pipelines/stable_diffusion/test_onnx_stable_diffusion_inpaint_legacy.py
@@ -70,6 +70,8 @@ def test_inference(self):
         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,
         )