Fix FreeU tests (huggingface#7540)

update
EMinsight · Apr 2, 2024 · 5d21d4a · 5d21d4a
1 parent 73ba810
commit 5d21d4a
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/tests/pipelines/test_pipelines_common.py b/tests/pipelines/test_pipelines_common.py
@@ -133,11 +133,15 @@ def test_freeu_enabled(self):
 
         inputs = self.get_dummy_inputs(torch_device)
         inputs["return_dict"] = False
+        inputs["output_type"] = "np"
+
         output = pipe(**inputs)[0]
 
         pipe.enable_freeu(s1=0.9, s2=0.2, b1=1.2, b2=1.4)
         inputs = self.get_dummy_inputs(torch_device)
         inputs["return_dict"] = False
+        inputs["output_type"] = "np"
+
         output_freeu = pipe(**inputs)[0]
 
         assert not np.allclose(
@@ -152,6 +156,8 @@ def test_freeu_disabled(self):
 
         inputs = self.get_dummy_inputs(torch_device)
         inputs["return_dict"] = False
+        inputs["output_type"] = "np"
+
         output = pipe(**inputs)[0]
 
         pipe.enable_freeu(s1=0.9, s2=0.2, b1=1.2, b2=1.4)
@@ -164,6 +170,8 @@ def test_freeu_disabled(self):
 
         inputs = self.get_dummy_inputs(torch_device)
         inputs["return_dict"] = False
+        inputs["output_type"] = "np"
+
         output_no_freeu = pipe(**inputs)[0]
         assert np.allclose(
             output, output_no_freeu, atol=1e-2