Skip to content

Commit 2b4f849

Browse files
[PipelineTesterMixin] Handle non-image outputs for attn slicing test (huggingface#2504)
* [PipelineTesterMixin] Handle non-image outputs for batch/sinle inference test * style --------- Co-authored-by: William Berman <WLBberman@gmail.com>
1 parent e4c356d commit 2b4f849

1 file changed

Lines changed: 5 additions & 2 deletions

File tree

‎tests/test_pipelines_common.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -450,7 +450,9 @@ def test_to_device(self):
450450
def test_attention_slicing_forward_pass(self):
451451
self._test_attention_slicing_forward_pass()
452452

453-
def _test_attention_slicing_forward_pass(self, test_max_difference=True, expected_max_diff=1e-3):
453+
def _test_attention_slicing_forward_pass(
454+
self, test_max_difference=True, test_mean_pixel_difference=True, expected_max_diff=1e-3
455+
):
454456
if not self.test_attention_slicing:
455457
return
456458

@@ -474,7 +476,8 @@ def _test_attention_slicing_forward_pass(self, test_max_difference=True, expecte
474476
max_diff = np.abs(output_with_slicing - output_without_slicing).max()
475477
self.assertLess(max_diff, expected_max_diff, "Attention slicing should not affect the inference results")
476478

477-
assert_mean_pixel_difference(output_with_slicing[0], output_without_slicing[0])
479+
if test_mean_pixel_difference:
480+
assert_mean_pixel_difference(output_with_slicing[0], output_without_slicing[0])
478481

479482
@unittest.skipIf(
480483
torch_device != "cuda" or not is_accelerate_available() or is_accelerate_version("<", "0.14.0"),

0 commit comments

Comments
 (0)